[llvm] [SLP] Add store-to-load forwarding conflict cost for widened store chains (PR #199606)

[Alexey Bataev via llvm-commits]

================
@@ -28885,6 +28930,210 @@ bool SLPVectorizerPass::runImpl(Function &F, ScalarEvolution *SE_,
   return Changed;
 }
 
+/// Returns the constant loop-carried byte stride of \p Ptr in \p L, i.e. the
+/// step of its affine SCEV recurrence, or std::nullopt when \p Ptr is not a
+/// simple affine recurrence in \p L with a constant step. This is the same
+/// notion as LoopAccessAnalysis's CommonStride; it is computed here directly
+/// from ScalarEvolution rather than via getPtrStride, whose no-wrap versioning
+/// and PredicatedScalarEvolution machinery is meant for legality, not costing.
+static std::optional<int64_t>
+getConstantLoopStrideInBytes(Value *Ptr, ScalarEvolution &SE, const Loop *L) {
+  const auto *AR = dyn_cast<SCEVAddRecExpr>(SE.getSCEV(Ptr));
+  if (!AR || AR->getLoop() != L)
+    return std::nullopt;
+  const auto *Step = dyn_cast<SCEVConstant>(AR->getStepRecurrence(SE));
+  if (!Step)
+    return std::nullopt;
+  const APInt &StepVal = Step->getAPInt();
+  if (StepVal.getSignificantBits() > 64)
+    return std::nullopt;
+  return StepVal.getSExtValue();
+}
+
+bool BoUpSLP::findStoreLoadForwardingConflict(
+    StoreInst *BaseStore, unsigned VF, LoadInst *OnlyLoad,
+    std::optional<uint64_t> LoadSizeOverride) {
+  assert(BaseStore && "Expected a valid base store");
+
+  Type *ValueTy = BaseStore->getValueOperand()->getType();
+  TypeSize StoreSize = DL->getTypeStoreSize(ValueTy);
+  if (StoreSize.isScalable())
+    return false;
+  uint64_t ElementSize = StoreSize.getFixedValue();
+  if (ElementSize == 0)
+    return false;
+
+  // Store-to-load forwarding hazards are a loop-carried concern.
+  const Loop *StoreL = LI->getLoopFor(BaseStore->getParent());
+  if (!StoreL)
+    return false;
+
+  uint64_t VectorStoreBytes = VF * ElementSize;
+  LLVM_DEBUG(dbgs() << "SLP: STLF check: VF=" << VF
+                    << " ElementSize=" << ElementSize
+                    << " VectorStoreBytes=" << VectorStoreBytes << "\n");
+
+  // Loop-carried byte stride of the store. A conflict is only a real hazard if
+  // a future iteration's load re-reads the bytes this store wrote, which is a
+  // property of the stride (see the per-load check below).
+  std::optional<int64_t> StoreStride =
+      getConstantLoopStrideInBytes(BaseStore->getPointerOperand(), *SE, StoreL);
+
+  // A store-to-load forwarding hazard can involve any load in the loop that
+  // reads the widened store's base, not only loads that became SLP tree nodes:
+  // a conflicting load may feed a scalar store, sit below a gather/splat leaf,
+  // or be vectorized in a different tree. Enumerate every simple load in the
+  // store's loop that shares the store base. The widened width below is only
+  // visible for loads in the current tree; loads vectorized by other trees are
+  // modeled at scalar width.
+  Value *StoreBase = getUnderlyingObject(BaseStore->getPointerOperand());
+  const auto CandidateLoads = [&] {
+    SmallPtrSet<LoadInst *, 8> Loads;
+    if (OnlyLoad) {
+      if (OnlyLoad->isSimple() &&
+          getUnderlyingObject(OnlyLoad->getPointerOperand()) == StoreBase)
+        Loads.insert(OnlyLoad);
+      return Loads;
+    }
+    for (BasicBlock *BB : StoreL->blocks())
+      for (Instruction &I : *BB)
+        if (auto *LoadI = dyn_cast<LoadInst>(&I))
+          if (LoadI->isSimple() &&
+              getUnderlyingObject(LoadI->getPointerOperand()) == StoreBase)
+            Loads.insert(LoadI);
+    return Loads;
+  }();
+
+  if (CandidateLoads.empty())
+    return false;
+
+  // For each candidate load, the widened chain becomes one wide store at the
+  // base; check whether the load straddles two such wide stores.
+  for (LoadInst *LoadI : CandidateLoads) {
+    // Only loads in the store's loop share its loop-carried dependence.
+    if (LI->getLoopFor(LoadI->getParent()) != StoreL)
+      continue;
+    const TreeEntry *WidenedLoadEntry = nullptr;
+    for (const TreeEntry *LTE : getTreeEntries(LoadI)) {
+      if (LTE->isGather() || DeletedNodes.contains(LTE) ||
+          TransformedToGatherNodes.contains(LTE))
+        continue;
+      if (LTE->hasState() && LTE->State == TreeEntry::Vectorize &&
+          LTE->getOpcode() == Instruction::Load) {
+        WidenedLoadEntry = LTE;
+        break;
+      }
+    }
+    std::optional<int64_t> Diff =
+        getPointersDiff(ValueTy, BaseStore->getPointerOperand(),
+                        LoadI->getType(), LoadI->getPointerOperand(), *DL, *SE,
+                        /*StrictCheck=*/true, /*CheckType=*/false);
+    if (!Diff || *Diff >= 0)
+      continue;
+    // Negating the minimum signed value is undefined.
+    if (*Diff == std::numeric_limits<int64_t>::min())
+      continue;
+
+    uint64_t Distance = static_cast<uint64_t>(-*Diff) * ElementSize;
+    LLVM_DEBUG(dbgs() << "SLP: STLF: load=" << *LoadI << " distance="
+                      << Distance << " bytes from chain base\n");
+
+    // A widened (regularly vectorized) load accesses the whole vector at once,
+    // so its effective width is the number of emitted lanes * element size, not
+    // one element. Such a wide load can straddle two wide stores even when
+    // perfectly aligned, which the misalignment-only test would miss. Use the
+    // count of distinct scalars actually loaded from memory (not the reuse-
+    // inflated vector factor) for the emitted load width.
+    TypeSize LoadTypeSize = DL->getTypeStoreSize(LoadI->getType());
+    uint64_t LoadElementSize =
+        LoadTypeSize.isScalable() ? 0 : LoadTypeSize.getFixedValue();
+    if (LoadSizeOverride)
+      LoadElementSize = *LoadSizeOverride;
+    else if (WidenedLoadEntry)
+      LoadElementSize *= WidenedLoadEntry->Scalars.size();
+    // A conflict is only a real hazard if a future iteration's load actually
+    // re-reads the bytes this store wrote. With a common positive loop-carried
+    // stride S, equal for the load and the store, the store's bytes are re-read
+    // iff there is an integer k >= 1 with
+    //   Distance - LoadElementSize < k * S < Distance + VectorStoreBytes.
+    // If no such k exists the accesses are strided-independent, so there is no
+    // forwarding hazard. When the stride is unknown, non-positive, or differs
+    // between load and store, fall back to the conservative check below.
+    std::optional<int64_t> LoadStride =
+        getConstantLoopStrideInBytes(LoadI->getPointerOperand(), *SE, StoreL);
+    if (StoreStride && LoadStride && *StoreStride == *LoadStride &&
----------------
alexey-bataev wrote:

```suggestion
    // The stride argument covers only re-reads of past iterations' stores;
    // a load wide enough to reach the current store must still go through
    // the conflict predicate.
    if (StoreStride && LoadStride && *StoreStride == *LoadStride &&
```

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