[llvm] [SLP] Avoid seeding related affine loop address computations (PR #226220)

[Tim Besard via llvm-commits]

================
@@ -1003,6 +1007,101 @@ bool isOnceUsedSeed(const Instruction *I) {
       I);
 }
 
+/// Returns true if \p Ptr is only used by scalar loads and stores, directly or
+/// through getelementptrs with constant offsets.
+static bool onlyFeedsScalarAccesses(Value *Ptr, bool ReVec,
+                                    unsigned Depth = 0) {
+  constexpr unsigned MaxDepth = 2;
+  if (Ptr->use_empty() || Ptr->hasNUsesOrMore(UsesLimit))
+    return false;
+  return all_of(Ptr->users(), [&](User *U) {
+    if (getValueType(U, ReVec)->isVectorTy())
+      return false;
+    Value *Op;
+    if (isa<LoadInst>(U) ||
+        (match(U, m_Store(m_Value(Op), m_Specific(Ptr))) && Op != Ptr))
+      return true;
+    if (auto *GEP = dyn_cast<GetElementPtrInst>(U))
+      return Depth < MaxDepth && GEP->getPointerOperand() == Ptr &&
+             GEP->hasAllConstantIndices() &&
+             onlyFeedsScalarAccesses(GEP, ReVec, Depth + 1);
+    return false;
+  });
+}
+
+/// Collects in \p GEPs the getelementptrs that the in-loop index computations
+/// \p VL end at, if they are only used by scalar accesses. Returns the loop
+/// containing all of them, or nullptr otherwise.
+static const Loop *
+collectScalarAccessGEPs(ArrayRef<Value *> VL, const LoopInfo &LI, bool ReVec,
+                        SmallVectorImpl<GetElementPtrInst *> &GEPs) {
+  constexpr unsigned MaxIndexChainLength = 3;
+  const Loop *L = nullptr;
+  for (Value *V : VL) {
+    Value *Cur = V;
+    GetElementPtrInst *GEP = nullptr;
+    for ([[maybe_unused]] unsigned _ : seq<unsigned>(MaxIndexChainLength)) {
+      if (!isa<BinaryOperator, CastInst>(Cur) || !Cur->hasOneUse())
+        return nullptr;
+      User *U = Cur->user_back();
+      if ((GEP = dyn_cast<GetElementPtrInst>(U)))
+        break;
+      Cur = U;
+    }
+    if (!GEP || GEP->getPointerOperand() == Cur ||
+        !onlyFeedsScalarAccesses(GEP, ReVec))
+      return nullptr;
+    // LSR only removes the arithmetic computed in the loop itself.
+    const Loop *GEPLoop = LI.getLoopFor(GEP->getParent());
+    if (!GEPLoop || (L && L != GEPLoop) ||
+        LI.getLoopFor(cast<Instruction>(V)->getParent()) != GEPLoop)
+      return nullptr;
+    L = GEPLoop;
----------------
maleadt wrote:

OK. `vectorizeOnceUsedSeeds` now only checks that the seed's index chain ends at a GEP in GEPs, and skips it if so. The loop, affinity and scalar-use checks are only in `isStrengthReducibleIndexBundle`, used by `vectorizeGEPIndices`.

Note that this skips any once-used seed feeding a GEP candidate of the same block, not only affine in-loop ones. Is this what you meant?

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