[llvm] [LAA] Allow vectorizing `A[NonZeroNonConstantStride*I] += 1` (PR #186262)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -969,54 +969,56 @@ class AccessAnalysis {
 
 } // end anonymous namespace
 
-std::optional<int64_t>
-llvm::getStrideFromAddRec(const SCEVAddRecExpr *AR, const Loop *Lp,
-                          Type *AccessTy, Value *Ptr,
-                          PredicatedScalarEvolution &PSE) {
+const SCEV *llvm::getStrideFromAddRec(const SCEVAddRecExpr *AR, const Loop *Lp,
+                                      Type *AccessTy, Value *Ptr,
+                                      PredicatedScalarEvolution &PSE) {
   if (isa<ScalableVectorType>(AccessTy)) {
     LLVM_DEBUG(dbgs() << "LAA: Bad stride - Scalable object: " << *AccessTy
                       << "\n");
-    return std::nullopt;
+    return nullptr;
   }
 
-  // The access function must stride over the innermost loop.
-  if (Lp != AR->getLoop()) {
+  auto BadStride = [&](auto Str) {
     LLVM_DEBUG({
-      dbgs() << "LAA: Bad stride - Not striding over innermost loop ";
+      dbgs() << "LAA: Bad stride - " << Str << " ";
       if (Ptr)
         dbgs() << *Ptr << " ";
 
       dbgs() << "SCEV: " << *AR << "\n";
     });
-    return std::nullopt;
-  }
+    return nullptr;
+  };
+
+  // The access function must stride over the innermost loop.
+  if (Lp != AR->getLoop())
+    return BadStride("Not striding over innermost loop");
+
+  // Check the step is loop invariant.
+  if (!AR->isAffine())
+    return BadStride("Step is varying");
 
-  // Check the step is constant.
   const SCEV *Step = AR->getStepRecurrence(*PSE.getSE());
 
-  // Calculate the pointer stride and check if it is constant.
-  const APInt *APStepVal;
-  if (!match(Step, m_scev_APInt(APStepVal))) {
-    LLVM_DEBUG({
-      dbgs() << "LAA: Bad stride - Not a constant strided ";
-      if (Ptr)
-        dbgs() << *Ptr << " ";
-      dbgs() << "SCEV: " << *AR << "\n";
-    });
-    return std::nullopt;
-  }
+  auto *SE = PSE.getSE();
+  const SCEV *AbsStep = SE->getAbsExpr(Step, false);
 
-  const auto &DL = Lp->getHeader()->getDataLayout();
-  TypeSize AllocSize = DL.getTypeAllocSize(AccessTy);
-  int64_t Size = AllocSize.getFixedValue();
+  const SCEV *TypeSizeScev = SE->getSizeOfExpr(
+      Step->getType(), SE->getDataLayout().getTypeAllocSize(AccessTy));
 
-  // Huge step value - give up.
-  std::optional<int64_t> StepVal = APStepVal->trySExtValue();
-  if (!StepVal)
-    return std::nullopt;
+  if (!SE->getURemExpr(AbsStep, TypeSizeScev)->isZero())
----------------
artagnon wrote:

I think we're allowed to add predicates to PSE here, which this function would facilitate:

```cpp
  LLVM_ABI bool
  isKnownMultipleOf(const SCEV *S, uint64_t M,
                    SmallVectorImpl<const SCEVPredicate *> &Assumptions);
```

https://github.com/llvm/llvm-project/pull/186262