Drop now unnecessary include

We now rely on the actual consequtive memops widening in a single place (both
original and speculated/mv'ed).

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -43,6 +43,11 @@ using namespace llvm;
 using namespace VPlanPatternMatch;
 using namespace SCEVPatternMatch;
 
+static cl::opt<bool> EnableVPlanBasedStrideMV(
+    "enable-vplan-based-stride-mv", cl::init(false), cl::Hidden,
+    cl::desc("Perform stride multiversioning directly on VPlan instead of in "
+             "LoopAccessAnalysis."));
+
 /// If the pointer operand \p Addr of a memory access is an affine AddRec
 /// w.r.t. \p L with a constant stride, return the stride in units of
 /// \p AccessTy. Otherwise return std::nullopt.
@@ -5454,6 +5459,11 @@ void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range,
         });
   }
 
+  if (EnableVPlanBasedStrideMV &&
+      !CostCtx.L->getHeader()->getParent()->hasOptSize())
+    RUN_VPLAN_PASS(VPlanTransforms::multiversionForUnitStridedMemOps, Plan,
+                   CostCtx, RecipeBuilder, Range, MemOps);
+
   // Widen unit-stride consecutive accesses, matching the legacy CM. Both
   // forward (stride +1) and reverse (stride -1) accesses are handled.
   VPlanTransforms::runPass(
@@ -5521,6 +5531,191 @@ void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range,
                            });
 }
 
+void VPlanTransforms::multiversionForUnitStridedMemOps(
+    VPlan &Plan, VPCostContext &CostCtx, VPRecipeBuilder &RecipeBuilder,
+    VFRange &Range, SmallVectorImpl<VPInstruction *> &MemOps) {
+  ScalarEvolution *SE = CostCtx.PSE.getSE();
+  PredicatedScalarEvolution StrideMVPSE(*SE, const_cast<Loop &>(*CostCtx.L));
+  SCEVUnionPredicate StridePredicates({}, *SE);
+
+  for (VPInstruction *VPI : MemOps) {
+    auto *PtrOp = VPI->getOpcode() == Instruction::Load ? 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 = VPI->getOpcode() == Instruction::Load
+                         ? VPI->getScalarType()
+                         : VPI->getOperand(0)->getScalarType();
+
+    const auto *TypeSize = cast<SCEVConstant>(SE->getSizeOfExpr(
+        Stride->getType(), SE->getDataLayout().getTypeAllocSize(ScalarTy)));
+
+    if (isa<SCEVConstant>(Stride))
+      continue;
+
+    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;
+    }
+
+    Value *StrideVal = cast<SCEVUnknown>(ToMultiVersion)->getValue();
+
+    const SCEVPredicate *NewPred =
+        SE->getComparePredicate(CmpInst::ICMP_EQ, ToMultiVersion, MVConst);
+
+    // Check if new predicate implies that backedge is never taken. If so, there
+    // is no reason to multiversion for it.
+    auto *PredicatedMaxBTC = SE->rewriteUsingPredicate(
+        SE->getSymbolicMaxBackedgeTakenCount(CostCtx.L), CostCtx.L,
+        StridePredicates.getUnionWith(NewPred, *SE)
+            .getUnionWith(&CostCtx.PSE.getPredicate(), *SE));
+
+    if (LoopVectorizationPlanner::getDecisionAndClampRange(
+            [&](ElementCount VF) {
+              return SE->isKnownPredicate(
+                  ICmpInst::ICMP_ULT, PredicatedMaxBTC,
+                  SE->getConstant(PredicatedMaxBTC->getType(),
+                                  VF.isScalable() ? 1
+                                                  : VF.getFixedValue() - 1));
+            },
+            Range))
+      continue;
+
+    if (VPI->getMask()) {
+      Instruction *I = VPI->getUnderlyingInstr();
+      bool IsLoad = VPI->getOpcode() == Instruction::Load;
+      // No reason to speculate unitstrideness if that won't improve vector code
+      // as we'd pay the price of not taking vector loop if the runtime
+      // condition is false for no benefits.
+      //
+      // We perform this clamping here because we don't want to
+      // `getDecisionAndClampRange` too early - here is better because we have
+      // another clamping right above.
+      if (!LoopVectorizationPlanner::getDecisionAndClampRange(
+              [&](ElementCount VF) -> bool {
+                return CostCtx.Config.isLegalMaskedLoadOrStore(
+                    IsLoad, ScalarTy, getLoadStoreAlignment(I),
+                    getLoadStoreAddressSpace(I));
+              },
+              Range))
+        continue;
+    }
+
+    StridePredicates = StridePredicates.getUnionWith(NewPred, *SE);
+
+    auto ReplaceMVUses = [&](Value *V) {
+      VPValue *From = Plan.getLiveIn(V);
+      if (!From)
+        return;
+      VPValue *To = Plan.getConstantInt(
+          From->getScalarType(),
+          cast<SCEVConstant>(MVConst)->getAPInt().getLimitedValue());
+
+      // 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.getVectorLoopRegion()->getSinglePredecessor();
+      });
+    };
+
+    ReplaceMVUses(StrideVal);
+    for (auto *U : StrideVal->users())
+      if (isa<SExtInst, ZExtInst, TruncInst>(U))
+        ReplaceMVUses(U);
+  }
+
+  if (StridePredicates.isAlwaysTrue())
+    return;
+
+  VPBasicBlock *Entry = Plan.getEntry();
+  VPBuilder Builder(Entry);
+  DebugLoc DL = cast<VPIRBasicBlock>(Entry)
+                    ->getIRBasicBlock()
+                    ->getTerminator()
+                    ->getDebugLoc();
+  VPSCEVExpander Expander(Builder, *SE, DL);
+  VPValue *Pred = Expander.tryToExpandPredicate(&StridePredicates);
+  assert(Pred && "Must be expandable!");
+
+  auto *StridesCheckBB = Plan.createVPBasicBlock("strides.check");
+  VPBasicBlock *ScalarPH = Plan.getScalarPreheader();
+  VPBlockUtils::insertBlockBefore(StridesCheckBB, Plan.getVectorPreheader());
+  VPBlockUtils::connectBlocks(StridesCheckBB, ScalarPH);
+  // SCEVExpander/VPSCEVExpander::expandCodeForPredicate negate the condition,
+  // so scalar preheader should be the first successor.
+  std::swap(StridesCheckBB->getSuccessors()[0],
+            StridesCheckBB->getSuccessors()[1]);
+  Builder.setInsertPoint(StridesCheckBB);
+  Builder.createNaryOp(VPInstruction::BranchOnCond, Pred);
+
+  for (VPRecipeBase &R : ScalarPH->phis()) {
+    auto &Phi = cast<VPPhi>(R);
+    Phi.addIncoming(Phi.getIncomingValueForBlock(Entry));
+  }
+
+  for (auto &R : make_early_inc_range(*Entry)) {
+    auto *ExpandSCEV = dyn_cast<VPExpandSCEVRecipe>(&R);
+    if (!ExpandSCEV)
+      continue;
+
+    const SCEV *S = ExpandSCEV->getSCEV();
+    Builder.setInsertPoint(ExpandSCEV);
+    const SCEV *NewS =
+        SE->rewriteUsingPredicate(S, CostCtx.L, StridePredicates);
+    if (NewS == S)
+      return;
+    auto *NewR = Builder.createExpandSCEV(NewS);
+    ExpandSCEV->replaceAllUsesWith(NewR);
+
+    // If this recipe is a trip count then we need to reset it explicitly.
+    if (ExpandSCEV == Plan.getTripCount())
+      Plan.resetTripCount(NewR);
+
+    ExpandSCEV->eraseFromParent();
+  }
+}
+
 void VPlanTransforms::makeScalarizationDecisions(VPlan &Plan, VFRange &Range) {
   if (LoopVectorizationPlanner::getDecisionAndClampRange(
           [&](ElementCount VF) { return VF.isScalar(); }, Range))

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -587,6 +587,12 @@ struct VPlanTransforms {
                                          VPRecipeBuilder &RecipeBuilder,
                                          VPCostContext &CostCtx);
 
+  /// \p MemOps must be updated to contain ones that haven't been processed by
+  /// the pass.
+  static void multiversionForUnitStridedMemOps(
+      VPlan &Plan, VPCostContext &CostCtx, VPRecipeBuilder &RecipeBuilder,
+      VFRange &Range, SmallVectorImpl<VPInstruction *> &MemOps);
+
   /// Make VPlan-based scalarization decision prior to delegating to the ones
   /// made by the legacy CM. Only transforms "usesFirstLaneOnly` def-use chains
   /// enabled by prior widening of consecutive memory operations for now.

llvm/lib/Transforms/Vectorize/VPlanUtils.cpp

@@ -1005,6 +1005,40 @@ VPValue *VPSCEVExpander::tryToExpand(const SCEV *S) {
   }
 }
 
+VPValue *VPSCEVExpander::tryToExpandPredicate(const SCEVPredicate *Pred) {
+  if (auto *CmpPred = dyn_cast<SCEVComparePredicate>(Pred)) {
+    VPValue *LHS = tryToExpand(CmpPred->getLHS());
+    if (!LHS)
+      return nullptr;
+    VPValue *RHS = tryToExpand(CmpPred->getRHS());
+    if (!RHS)
+      return nullptr;
+    // Match SCEVExpander::expandComparePredicate behavior:
+    auto InvPred = CmpInst::getInversePredicate(CmpPred->getPredicate());
+    return Builder.createICmp(InvPred, LHS, RHS, DL);
+  }
+
+  if (auto *UnionPred = dyn_cast<SCEVUnionPredicate>(Pred)) {
+    SmallVector<VPValue *, 4> Checks;
+    for (const auto *Pred : UnionPred->getPredicates()) {
+      VPValue *Check = tryToExpandPredicate(Pred);
+      if (!Check)
+        return nullptr;
+      Checks.push_back(Check);
+    }
+
+    if (Checks.empty())
+      return Builder.getPlan().getFalse();
+
+    VPValue *Result = Checks[0];
+    for (VPValue *Check : drop_begin(Checks))
+      Result = Builder.createOr(Result, Check, DL);
+    return Result;
+  }
+
+  llvm_unreachable("Unhandled predicate");
+}
+
 bool vputils::isDeadRecipe(VPRecipeBase &R) {
   // Do remove conditional assume instructions as their conditions may be
   // flattened.

llvm/lib/Transforms/Vectorize/VPlanUtils.h

@@ -264,6 +264,11 @@ public:
   /// Try to expand \p S into recipes and live-ins using the builder. Returns
   /// nullptr if \p S cannot be expanded yet.
   VPValue *tryToExpand(const SCEV *S);
+
+  /// Try to expand \p Pred into recipes and live-ins using the builder. Returns
+  /// nullptr if \p Pred cannot be expanded (not expected to happen for the
+  /// predicates we currently generate).
+  VPValue *tryToExpandPredicate(const SCEVPredicate *Pred);
 };
 //===----------------------------------------------------------------------===//
 // Utilities for modifying predecessors and successors of VPlan blocks.