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

[Luke Lau via llvm-commits]

================
@@ -5493,6 +5502,193 @@ void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range,
                            });
 }
 
+void VPlanTransforms::multiversionForUnitStridedMemOps(
+    VPlan &Plan, VPCostContext &CostCtx, VFRange &Range,
+    ArrayRef<VPInstruction *> MemOps) {
+  ScalarEvolution *SE = CostCtx.PSE.getSE();
+  SCEVUnionPredicate StridePredicates({}, *SE);
+
+  for (VPInstruction *VPI : MemOps) {
+    VPValue *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();
+
+    if (VPI->getMask()) {
+      Instruction *I = VPI->getUnderlyingInstr();
+      bool IsLoad = VPI->getOpcode() == Instruction::Load;
+      // 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;
+
+    // Corner case probably, `stride == [sz]ext i1 %val`. For stores, LAA would
+    // probably introduce run-time checks to avoid stores to a potentially loop
+    // invariant address and for loads a more effective strategy would be to
+    // have a branch inside loop body to choose between unit-strided/broadcasted
+    // load, if that ever becomes important.
+    //
+    // As a bonus, we wouldn't need to worry about `sext(i1 1)` being negative.
+    if (ToMultiVersion->getType()->isIntegerTy(1))
+      continue;
----------------
lukel97 wrote:

IIUC this is new functionality since this isn't something that LAA guarded against previously, so keeping the behaviour in sync helps reduce diffs when we switch to vplan multiversioning by default.

I'm not opposed to skipping multi-versioning for i1 strides in the long run but I'd rather see that split out into another PR so we can reason about it separately

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