[llvm] [VPlan] Implement VPlan-based unit-strideness speculation (PR #182595)

[Mel Chen via llvm-commits]

================
@@ -5668,6 +5677,198 @@ void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range,
                            });
 }
 
+/// Replace the stride and all integral casts derived from it. Some casts might
+/// be defined outside VPlan and modeled as live-ins only, so process users at
+/// the LLVM IR level to discover them.
+static void replaceUsesInVectorLoopThroughCastsRecursively(VPlan &Plan,
+                                                           ScalarEvolution &SE,
+                                                           Value *V,
+                                                           const SCEV *ToSCEV) {
+  VPValue *From = Plan.getLiveIn(V);
+  if (From) {
+    assert(From->getScalarType() == ToSCEV->getType() &&
+           "Wrong type for ToSCEV!");
+    VPValue *To = Plan.getConstantInt(cast<SCEVConstant>(ToSCEV)->getAPInt());
+
+    // Original scalar loop can still use `From`, make sure to only rewrite
+    // uses inside the vector loop that we guard with the checks.
+    From->replaceUsesWithIf(To, [&](VPUser &U, unsigned) {
+      auto *R = cast<VPRecipeBase>(&U);
+      return R->getRegion() || R->getParent() == Plan.getVectorPreheader();
+    });
+  }
+
+  for (User *U : V->users())
+    if (isa<SExtInst>(U))
+      replaceUsesInVectorLoopThroughCastsRecursively(
+          Plan, SE, U, SE.getSignExtendExpr(ToSCEV, U->getType()));
+    else if (isa<ZExtInst, TruncInst>(U))
+      replaceUsesInVectorLoopThroughCastsRecursively(
+          Plan, SE, U, SE.getTruncateOrZeroExtend(ToSCEV, U->getType()));
+}
+
+void VPlanTransforms::multiversionForUnitStridedMemOps(
+    VPlan &Plan, VPCostContext &CostCtx, VFRange &Range,
+    ArrayRef<VPInstruction *> MemOps) {
+  ScalarEvolution *SE = CostCtx.PSE.getSE();
+  SCEVUnionPredicate StridePredicates({}, *SE);
+
+  for (VPInstruction *VPI : MemOps) {
+    bool IsLoad = VPI->getOpcode() == Instruction::Load;
+    VPValue *PtrOp = IsLoad ? VPI->getOperand(0) : VPI->getOperand(1);
+
+    const SCEV *PtrSCEV =
+        vputils::getSCEVExprForVPValue(PtrOp, CostCtx.PSE, CostCtx.L);
+    const SCEV *Start, *Stride;
+
+    if (!match(PtrSCEV, m_scev_AffineAddRec(m_SCEV(Start), m_SCEV(Stride),
+                                            m_SpecificLoop(CostCtx.L))))
+      continue;
+
+    Type *ScalarTy =
+        IsLoad ? VPI->getScalarType() : VPI->getOperand(0)->getScalarType();
+
+    if (VPI->getMask()) {
+      Instruction *I = VPI->getUnderlyingInstr();
+      // Don't speculate unit-strideness if it won't result in any unit-strided
+      // loads, as we'd pay the price of not taking vector loop if the runtime
+      // condition is false for no benefits.
+      if (!CostCtx.Config.isLegalMaskedLoadOrStore(IsLoad, ScalarTy,
+                                                   getLoadStoreAlignment(I),
+                                                   getLoadStoreAddressSpace(I)))
+        continue;
+    }
+
+    if (isa<SCEVConstant>(Stride))
+      continue;
+
+    const auto *TypeSize = cast<SCEVConstant>(SE->getSizeOfExpr(
+        Stride->getType(), SE->getDataLayout().getTypeAllocSize(ScalarTy)));
+
+    const SCEVConstant *StrideConstantMultiplier;
+    const SCEV *StrideNonConstantMultiplier;
+
+    const SCEV *ToMultiVersion = Stride;
+    const SCEV *MVConst = TypeSize;
+    if (match(Stride, m_scev_c_Mul(m_SCEVConstant(StrideConstantMultiplier),
+                                   m_SCEV(StrideNonConstantMultiplier)))) {
+      if (TypeSize != StrideConstantMultiplier) {
+        // TODO: Support `TypeSize = N * StrideConstantMultiplier`,
+        // including negative `N`. For now, only process when they're equal,
+        // which matches the useful part of the legacy behavior that
+        // multiversiones GEP index for stride one.
+        continue;
+      }
+      ToMultiVersion = StrideNonConstantMultiplier;
+      MVConst = SE->getOne(ToMultiVersion->getType());
+    } else if (!TypeSize->isOne()) {
+      // Likewise - try to match legacy behavior.
+      continue;
+    }
+
+    while (auto *C = dyn_cast<SCEVIntegralCastExpr>(ToMultiVersion)) {
+      ToMultiVersion = C->getOperand();
+      MVConst = SE->getTruncateOrSignExtend(MVConst, ToMultiVersion->getType());
+    }
+
+    if (match(ToMultiVersion, m_scev_UndefOrPoison()))
+      continue;
+
+    if (!isa<SCEVUnknown>(ToMultiVersion)) {
+      // Match legacy behavior.
+      // If/when changed, make sure that explicit poison/undef in the defining
+      // expression doesn't cause any issues.
+      continue;
+    }
+
+    // This should probably use a rewrite using a union between
+    // PSE/StridePredicates, but I can't a craft a test as we're bailing out
+    // earlied via `isa<SCEVUnknown>`.
+    if (SE->isKnownPredicate(CmpInst::ICMP_NE, ToMultiVersion, MVConst))
+      continue;
+
+    Value *StrideVal = cast<SCEVUnknown>(ToMultiVersion)->getValue();
+
+    const SCEVPredicate *NewPred =
+        SE->getComparePredicate(CmpInst::ICMP_EQ, ToMultiVersion, MVConst);
----------------
Mel-Chen wrote:

In the future, is it possible that within the same vectorized loop, multiversioning might require different predicate assumptions? If so, which predicate should take precedence?

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