[LV] Model stride in VPWidenMemoryInstructionRecipe [nfc]

[reames]

Summary

We currently allow two strides 1 for consecutive accesses, and -1 for reverse consecutive access. In the near future, I want to add support for other constant strides, and getting the API worked through simplifies that work. This is admittedly somewhat debatable on its own merits as an NFC, but I figured I'd put it out and see what reviewers thought.

[fhahn]

Do you have any particular examples in mind? At the moment, strides > 1 would be handled as interleave group I think and maybe that would also work.

[reames]

In D148388#4287713, @fhahn wrote:

Do you have any particular examples in mind? At the moment, strides > 1 would be handled as interleave group I think and maybe that would also work.

Constant strides greater than 8 currently end up as masked loads or stores, and the addressing for a masked load vs a strided load is different. The latter only uses the first lane of the pointer vector.

We could also use it for strided stores (even less than stride 8). On a target with wide loads, but not masking, this would enable strided patterns entirely. On a target with wide memory ops and masking there can still be a performance difference between strided access and masking.

(For context, RISCV has a native strided load and store instruction.)

And all of the above is really just a building block to removing the stride==1 speculation and handling runtime strides. This is the case I actually care about, I just need to get the codegen part into acceptable shape first. (This shows up in spec2017 x264 in several cases.)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

if (!LoopVectorizationPlanner::getDecisionAndClampRange(willWiden, Range))
     return nullptr;
 
   VPValue *Mask = nullptr;
   if (Legal->isMaskRequired(I))
     Mask = createBlockInMask(I->getParent(), *Plan);
 
   // Determine if the pointer operand of the access is either consecutive or
   // reverse consecutive.
-  LoopVectorizationCostModel::InstWidening Decision =
-      CM.getWideningDecision(I, Range.Start);
-  bool Reverse = Decision == LoopVectorizationCostModel::CM_Widen_Reverse;
-  bool Consecutive =
-      Reverse || Decision == LoopVectorizationCostModel::CM_Widen;
+  const std::optional<int> Stride = [&]() -> std::optional<int> {
+    switch (CM.getWideningDecision(I, Range.Start)) {
+    default:
+      return std::nullopt;
+    case LoopVectorizationCostModel::CM_Widen_Reverse:
+      return -1;
+    case LoopVectorizationCostModel::CM_Widen:
+      return 1;
+    }
+  }();
 
   if (LoadInst *Load = dyn_cast<LoadInst>(I))
     return new VPWidenMemoryInstructionRecipe(*Load, Operands[0], Mask,
-                                              Consecutive, Reverse);
+                                              Stride);
 
   StoreInst *Store = cast<StoreInst>(I);
   return new VPWidenMemoryInstructionRecipe(*Store, Operands[1], Operands[0],
-                                            Mask, Consecutive, Reverse);
+                                            Mask, Stride);
 }
 
 /// Creates a VPWidenIntOrFpInductionRecpipe for \p Phi. If needed, it will also
 /// insert a recipe to expand the step for the induction recipe.
 static VPWidenIntOrFpInductionRecipe *createWidenInductionRecipes(
     PHINode *Phi, Instruction *PhiOrTrunc, VPValue *Start,
     const InductionDescriptor &IndDesc, LoopVectorizationCostModel &CM,
     VPlan &Plan, ScalarEvolution &SE, Loop &OrigLoop, VFRange &Range) {
void VPWidenMemoryInstructionRecipe::execute(VPTransformState &State) {
   assert((LI || SI) && "Invalid Load/Store instruction");
   assert((!SI || StoredValue) && "No stored value provided for widened store");
   assert((!LI || !StoredValue) && "Stored value provided for widened load");
 
   Type *ScalarDataTy = getLoadStoreType(&Ingredient);
 
   auto *DataTy = VectorType::get(ScalarDataTy, State.VF);
   const Align Alignment = getLoadStoreAlignment(&Ingredient);
-  bool CreateGatherScatter = !Consecutive;
+  bool CreateGatherScatter = !isConsecutive();
 
   auto &Builder = State.Builder;
   InnerLoopVectorizer::VectorParts BlockInMaskParts(State.UF);
   bool isMaskRequired = getMask();
   if (isMaskRequired)
     for (unsigned Part = 0; Part < State.UF; ++Part)
       BlockInMaskParts[Part] = State.get(getMask(), Part);
 
   const auto CreateVecPtr = [&](unsigned Part, Value *Ptr) -> Value * {
     // Calculate the pointer for the specific unroll-part.
     Value *PartPtr = nullptr;
 
     // Use i32 for the gep index type when the value is constant,
     // or query DataLayout for a more suitable index type otherwise.
     const DataLayout &DL =
         Builder.GetInsertBlock()->getModule()->getDataLayout();
-    Type *IndexTy = State.VF.isScalable() && (Reverse || Part > 0)
+    Type *IndexTy = State.VF.isScalable() && (isReverse() || Part > 0)
                         ? DL.getIndexType(ScalarDataTy->getPointerTo())
                         : Builder.getInt32Ty();
     bool InBounds = false;
     if (auto *gep = dyn_cast<GetElementPtrInst>(Ptr->stripPointerCasts()))
       InBounds = gep->isInBounds();
-    if (Reverse) {
+    if (isReverse()) {
       // If the address is consecutive but reversed, then the
       // wide store needs to start at the last vector element.
       // RunTimeVF =  VScale * VF.getKnownMinValue()
       // For fixed-width VScale is 1, then RunTimeVF = VF.getKnownMinValue()
       Value *RunTimeVF = getRuntimeVF(Builder, IndexTy, State.VF);
       // NumElt = -Part * RunTimeVF
       Value *NumElt =
           Builder.CreateMul(ConstantInt::get(IndexTy, -(int64_t)Part), RunTimeVF);
for (unsigned Part = 0; Part < State.UF; ++Part) {
       Instruction *NewSI = nullptr;
       Value *StoredVal = State.get(StoredValue, Part);
       if (CreateGatherScatter) {
         Value *MaskPart = isMaskRequired ? BlockInMaskParts[Part] : nullptr;
         Value *VectorGep = State.get(getAddr(), Part);
         NewSI = Builder.CreateMaskedScatter(StoredVal, VectorGep, Alignment,
                                             MaskPart);
       } else {
-        if (Reverse) {
+        if (isReverse()) {
           // If we store to reverse consecutive memory locations, then we need
           // to reverse the order of elements in the stored value.
           StoredVal = Builder.CreateVectorReverse(StoredVal, "reverse");
           // We don't want to update the value in the map as it might be used in
           // another expression. So don't call resetVectorValue(StoredVal).
         }
         auto *VecPtr =
             CreateVecPtr(Part, State.get(getAddr(), VPIteration(0, 0)));
if (CreateGatherScatter) {
             DataTy, VecPtr, Alignment, BlockInMaskParts[Part],
             PoisonValue::get(DataTy), "wide.masked.load");
       else
         NewLI =
             Builder.CreateAlignedLoad(DataTy, VecPtr, Alignment, "wide.load");
 
       // Add metadata to the load, but setVectorValue to the reverse shuffle.
       State.addMetadata(NewLI, LI);
-      if (Reverse)
+      if (isReverse())
         NewLI = Builder.CreateVectorReverse(NewLI, "reverse");
     }
 
     State.set(getVPSingleValue(), NewLI, Part);
   }
 }
 
 // Determine how to lower the scalar epilogue, which depends on 1) optimising

llvm/lib/Transforms/Vectorize/VPlan.h

 /// The recipe uses the following VPValues:
 /// - For load: Address, optional mask
 /// - For store: Address, stored value, optional mask
 /// TODO: We currently execute only per-part unless a specific instance is
 /// provided.
 class VPWidenMemoryInstructionRecipe : public VPRecipeBase {
   Instruction &Ingredient;
 
-  // Whether the loaded-from / stored-to addresses are consecutive.
-  bool Consecutive;
-
-  // Whether the consecutive loaded/stored addresses are in reverse order.
-  bool Reverse;
+  // The (constant) stride of the memory access as lowered.  Note that this is
+  // a lowering choice; there can be strided accesses in the source which we've
+  // decided to lower via gather/scatter and they will not have the stride set
+  // here.  Interesting special cases are:
+  // *  1 -- the loaded-from / stored-to addresses are consecutive.  Will be
+  //    lowered as a single wide access.
+  // * -1 -  the consecutive loaded/stored addresses are in reverse order.
+  //    Will be lowered as a wide access followed by a vector reverse.
+  std::optional<int> Stride;
 
   void setMask(VPValue *Mask) {
     if (!Mask)
       return;
     addOperand(Mask);
   }
 
   bool isMasked() const {
     return isStore() ? getNumOperands() == 3 : getNumOperands() == 2;
   }
 
 public:
   VPWidenMemoryInstructionRecipe(LoadInst &Load, VPValue *Addr, VPValue *Mask,
-                                 bool Consecutive, bool Reverse)
+                                 std::optional<int> Stride)
       : VPRecipeBase(VPDef::VPWidenMemoryInstructionSC, {Addr}),
-        Ingredient(Load), Consecutive(Consecutive), Reverse(Reverse) {
-    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
+        Ingredient(Load), Stride(Stride) {
     new VPValue(this, &Load);
     setMask(Mask);
   }
 
   VPWidenMemoryInstructionRecipe(StoreInst &Store, VPValue *Addr,
                                  VPValue *StoredValue, VPValue *Mask,
-                                 bool Consecutive, bool Reverse)
+                                 std::optional<int> Stride)
       : VPRecipeBase(VPDef::VPWidenMemoryInstructionSC, {Addr, StoredValue}),
-        Ingredient(Store), Consecutive(Consecutive), Reverse(Reverse) {
-    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
+        Ingredient(Store), Stride(Stride) {
     setMask(Mask);
   }
 
   VP_CLASSOF_IMPL(VPDef::VPWidenMemoryInstructionSC)
 
   /// Return the address accessed by this recipe.
   VPValue *getAddr() const {
     return getOperand(0); // Address is the 1st, mandatory operand.
public:
 
   /// Return the address accessed by this recipe.
   VPValue *getStoredValue() const {
     assert(isStore() && "Stored value only available for store instructions");
     return getOperand(1); // Stored value is the 2nd, mandatory operand.
   }
 
   // Return whether the loaded-from / stored-to addresses are consecutive.
-  bool isConsecutive() const { return Consecutive; }
+  bool isConsecutive() const {
+    return Stride && (*Stride == 1 || *Stride == -1);
+  }
 
   // Return whether the consecutive loaded/stored addresses are in reverse
   // order.
-  bool isReverse() const { return Reverse; }
+  bool isReverse() const { return Stride && *Stride == -1; }
 
   /// Generate the wide load/store.
   void execute(VPTransformState &State) override;
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (VPRecipeBase &Ingredient :
       } else {
         assert(isa<VPInstruction>(&Ingredient) &&
                "only VPInstructions expected here");
         assert(!isa<PHINode>(Inst) && "phis should be handled above");
         // Create VPWidenMemoryInstructionRecipe for loads and stores.
         if (LoadInst *Load = dyn_cast<LoadInst>(Inst)) {
           NewRecipe = new VPWidenMemoryInstructionRecipe(
               *Load, Plan->getOrAddVPValue(getLoadStorePointerOperand(Inst)),
-              nullptr /*Mask*/, false /*Consecutive*/, false /*Reverse*/);
+              nullptr /*Mask*/, std::nullopt /*stride*/);
         } else if (StoreInst *Store = dyn_cast<StoreInst>(Inst)) {
           NewRecipe = new VPWidenMemoryInstructionRecipe(
               *Store, Plan->getOrAddVPValue(getLoadStorePointerOperand(Inst)),
               Plan->getOrAddVPValue(Store->getValueOperand()), nullptr /*Mask*/,
-              false /*Consecutive*/, false /*Reverse*/);
+              std::nullopt /*stride*/);
         } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Inst)) {
           NewRecipe =
               new VPWidenGEPRecipe(GEP, Plan->mapToVPValues(GEP->operands()));
         } else if (CallInst *CI = dyn_cast<CallInst>(Inst)) {
           NewRecipe =
               new VPWidenCallRecipe(*CI, Plan->mapToVPValues(CI->args()),
                                     getVectorIntrinsicIDForCall(CI, &TLI));
         } else if (SelectInst *SI = dyn_cast<SelectInst>(Inst)) {

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

TEST(VPRecipeTest, CastVPWidenMemoryInstructionRecipeToVPUserAndVPDef) {
   LLVMContext C;
 
   IntegerType *Int32 = IntegerType::get(C, 32);
   PointerType *Int32Ptr = PointerType::get(Int32, 0);
   auto *Load =
       new LoadInst(Int32, UndefValue::get(Int32Ptr), "", false, Align(1));
   VPValue Addr;
   VPValue Mask;
-  VPWidenMemoryInstructionRecipe Recipe(*Load, &Addr, &Mask, true, false);
+  VPWidenMemoryInstructionRecipe Recipe(*Load, &Addr, &Mask, 1);
   EXPECT_TRUE(isa<VPUser>(&Recipe));
   VPRecipeBase *BaseR = &Recipe;
   EXPECT_TRUE(isa<VPUser>(BaseR));
   EXPECT_EQ(&Recipe, BaseR);
 
   VPValue *VPV = Recipe.getVPSingleValue();
   EXPECT_TRUE(isa<VPRecipeBase>(VPV->getDefiningRecipe()));
   EXPECT_EQ(&Recipe, VPV->getDefiningRecipe());
PointerType *Int32Ptr = PointerType::get(Int32, 0);
     EXPECT_FALSE(Recipe.mayReadOrWriteMemory());
   }
 
   {
     auto *Load =
         new LoadInst(Int32, UndefValue::get(Int32Ptr), "", false, Align(1));
     VPValue Addr;
     VPValue Mask;
-    VPWidenMemoryInstructionRecipe Recipe(*Load, &Addr, &Mask, true, false);
+    VPWidenMemoryInstructionRecipe Recipe(*Load, &Addr, &Mask, 1);
     EXPECT_FALSE(Recipe.mayHaveSideEffects());
     EXPECT_TRUE(Recipe.mayReadFromMemory());
     EXPECT_FALSE(Recipe.mayWriteToMemory());
     EXPECT_TRUE(Recipe.mayReadOrWriteMemory());
     delete Load;
   }
 
   {
     auto *Store = new StoreInst(UndefValue::get(Int32),
                                 UndefValue::get(Int32Ptr), false, Align(1));
     VPValue Addr;
     VPValue Mask;
     VPValue StoredV;
-    VPWidenMemoryInstructionRecipe Recipe(*Store, &Addr, &StoredV, &Mask, false,
-                                          false);
+    VPWidenMemoryInstructionRecipe Recipe(*Store, &Addr, &StoredV, &Mask,
+                                          std::nullopt);
     EXPECT_TRUE(Recipe.mayHaveSideEffects());
     EXPECT_FALSE(Recipe.mayReadFromMemory());
     EXPECT_TRUE(Recipe.mayWriteToMemory());
     EXPECT_TRUE(Recipe.mayReadOrWriteMemory());
     delete Store;
   }
 
   {