[llvm] [LoopAccessAnalysis] Analyze forked pointer dependences per alternative (PR #214216)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -2234,6 +2243,112 @@ MemoryDepChecker::getDependenceDistanceStrideAndSize(
                                       TypeByteSize, AIsWrite, BIsWrite);
 }
 
+std::optional<int64_t>
+MemoryDepChecker::getStrideFromSCEV(
+    const SCEV *S, Type *AccessTy,
+    SmallVectorImpl<const SCEVPredicate *> *Predicates) {
+  ScalarEvolution &SE = *PSE.getSE();
+  if (SE.isLoopInvariant(S, InnermostLoop))
+    return 0;
+  const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(S);
+  if (!AR)
+    return std::nullopt;
+  std::optional<int64_t> Stride =
+      getStrideFromAddRec(AR, InnermostLoop, AccessTy, /*Ptr=*/nullptr, PSE);
+  if (!Stride)
+    return std::nullopt;
+  if (isNoWrap(PSE, AR, /*Ptr=*/nullptr, AccessTy, InnermostLoop, *DT, Stride,
+               Predicates))
+    return Stride;
+  return std::nullopt;
+}
+
+std::optional<MemoryDepChecker::Dependence::DepType>
+MemoryDepChecker::getForkedDepType(const MemAccessInfo &A, Instruction *AInst,
+                                   const MemAccessInfo &B, Instruction *BInst) {
+  // Two reads are independent.
+  const auto &[APtr, AIsWrite] = A;
+  const auto &[BPtr, BIsWrite] = B;
+  if (!AIsWrite && !BIsWrite)
+    return Dependence::NoDep;
+
+  ScalarEvolution &SE = *PSE.getSE();
+
+  // A pointer may be a fork of multiple strided pointers, e.g. produced by a
+  // select of two pointers. SCEV cannot form a single AddRec for such a
+  // pointer, so the regular analysis in getDependenceDistanceStrideAndSize
+  // would report an IndirectUnsafe dependence. The runtime pointer checking
+  // handles such pointers by analyzing each fork alternative separately
+  // (findForkedSCEVs); do the same here by checking each pair of alternatives
+  // and aggregating the results.
+  SmallVector<PointerIntPair<const SCEV *, 1, bool>> Srcs;
+  SmallVector<PointerIntPair<const SCEV *, 1, bool>> Sinks;
+  findForkedSCEVs(&SE, InnermostLoop, APtr, Srcs, MaxForkedSCEVDepth);
+  findForkedSCEVs(&SE, InnermostLoop, BPtr, Sinks, MaxForkedSCEVDepth);
+
+  // If neither access pointer is a fork, fall back to the regular single-SCEV
+  // analysis in getDependenceDistanceStrideAndSize.
+  if (Srcs.size() == 1 && Sinks.size() == 1)
+    return std::nullopt;
+
+  // We can only analyze a fork if each alternative is loop-invariant or an
+  // affine AddRec; this is the same requirement as for generating runtime
+  // checks for forked pointers (see AccessAnalysis::createCheckForAccess).
+  auto IsLoopInvariantOrAR =
+      [&](const PointerIntPair<const SCEV *, 1, bool> &P) {
+        return SE.isLoopInvariant(P.getPointer(), InnermostLoop) ||
+               isa<SCEVAddRecExpr>(P.getPointer());
+      };
+  if (!all_of(Srcs, IsLoopInvariantOrAR) ||
+      !all_of(Sinks, IsLoopInvariantOrAR))
+    return Dependence::IndirectUnsafe;
+
+  // We cannot check pointers in different address spaces.
+  if (APtr->getType()->getPointerAddressSpace() !=
+      BPtr->getType()->getPointerAddressSpace())
+    return Dependence::Unknown;
+
+  Type *ATy = getLoadStoreType(AInst);
+  Type *BTy = getLoadStoreType(BInst);
+
+  SmallVector<const SCEVPredicate *> Predicates;
+  Dependence::DepType Result = Dependence::NoDep;
+  bool ResultNeedsRtCheck = false;
+  for (const auto &[SrcArm, _] : Srcs) {
+    for (const auto &[SinkArm, _] : Sinks) {
----------------
artagnon wrote:

Are we taking non-corresponding sources and sinks? We should check that only one of them forked?

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