[LoopVectorize][LAA] Hoist load in memory IV to allow vectorization

Adds VPScalarIVPromotionRecipe recipe to promote memory IV to scalar
IV.

Use SCEV mutilplied by VFxUF or EVL to multiply the SCEV between load
and store.

The code that this patch allows vectorization is like this:

while.body:
  %theFirst.addr.0112 = phi ptr [ %incdec.ptr9, %while.body ], [ %theFirst, %while.body.preheader ]
  %thePointer.0111 = phi ptr [ %incdec.ptr, %while.body ], [ %add.ptr.i, %while.body.preheader ]
  %1 = load i16, ptr %theFirst.addr.0112, align 2
  store i16 %1, ptr %thePointer.0111, align 2
  %incdec.ptr = getelementptr inbounds nuw i8, ptr %thePointer.0111, i64 2
  %2 = load i64, ptr %m_size_ptr, align 8
  %inc = add i64 %2, 1
  store i64 %inc, ptr %m_size_ptr, align 8
  %incdec.ptr9 = getelementptr inbounds nuw i8, ptr %theFirst.addr.0112, i64 2
  %cmp7.not = icmp eq ptr %incdec.ptr9, %theLast
  br i1 %cmp7.not, label %cleanup.loopexit, label %while.body

As you can see, %m_size_ptr is a loop invariant pointer and can be
promoted to a IV scalar and then execute vectorization.

Author: felipealmeida

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -2329,6 +2329,19 @@ MemoryDepChecker::getDependenceDistanceStrideAndSize(
   int64_t StrideBPtrInt = *StrideBPtr;
   LLVM_DEBUG(dbgs() << "LAA:  Src induction step: " << StrideAPtrInt
                     << " Sink induction step: " << StrideBPtrInt << "\n");
+
+  if (!StrideAPtrInt && !StrideBPtrInt && !(AIsWrite && BIsWrite) &&
+      (AIsWrite || BIsWrite) && !isa<UndefValue>(APtr) &&
+      InnermostLoop->isLoopInvariant(APtr) &&
+      InnermostLoop->isLoopInvariant(BPtr)) {
+    LoadInst *L = dyn_cast<LoadInst>(AIsWrite ? BInst : AInst);
+    if (InnermostLoop->isLoopInvariant(L->getPointerOperand()))
+      if (L && isInvariantLoadHoistable(L, SE, nullptr, nullptr, nullptr))
+        ShouldRetryWithRuntimeChecks = true;
+
+    return MemoryDepChecker::Dependence::Unknown;
+  }
+
   // At least Src or Sink are loop invariant and the other is strided or
   // invariant. We can generate a runtime check to disambiguate the accesses.
   if (!StrideAPtrInt || !StrideBPtrInt)
@@ -2942,9 +2955,15 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, const LoopInfo *LI,
     // See if there is an unsafe dependency between a load to a uniform address and
     // store to the same uniform address.
     if (UniformStores.contains(Ptr)) {
-      LLVM_DEBUG(dbgs() << "LAA: Found an unsafe dependency between a uniform "
-                           "load and uniform store to the same address!\n");
-      HasLoadStoreDependenceInvolvingLoopInvariantAddress = true;
+      auto &SE = *PSE->getSE();
+      if (TheLoop->isLoopInvariant(LD->getPointerOperand()) &&
+          !getDepChecker().isInvariantLoadHoistable(LD, SE, nullptr, nullptr,
+                                                    nullptr)) {
+        LLVM_DEBUG(
+            dbgs() << "LAA: Found an unsafe dependency between a uniform "
+                      "load and uniform store to the same address!\n");
+        HasLoadStoreDependenceInvolvingLoopInvariantAddress = true;
+      }
     }
 
     MemoryLocation Loc = MemoryLocation::get(LD);

llvm/lib/Transforms/Scalar/LoopDistribute.cpp

@@ -680,12 +680,14 @@ class LoopDistributeForLoop {
 
     // Currently, we only distribute to isolate the part of the loop with
     // dependence cycles to enable partial vectorization.
-    if (LAI->canVectorizeMemory())
+    if (!LAI->hasLoadStoreDependenceInvolvingLoopInvariantAddress() &&
+        LAI->canVectorizeMemory())
       return fail("MemOpsCanBeVectorized",
                   "memory operations are safe for vectorization");
 
     auto *Dependences = LAI->getDepChecker().getDependences();
-    if (!Dependences || Dependences->empty())
+    if (!LAI->hasLoadStoreDependenceInvolvingLoopInvariantAddress() &&
+        (!Dependences || Dependences->empty()))
       return fail("NoUnsafeDeps", "no unsafe dependences to isolate");
 
     LLVM_DEBUG(dbgs() << "LDist: Found a candidate loop: "

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

@@ -630,6 +630,8 @@ class LoopVectorizationPlanner {
                                   VPRecipeBuilder &RecipeBuilder,
                                   ElementCount MinVF);
 
+  void adjustScalarIVPromotions(VPlanPtr &Plan);
+
   /// Attach the runtime checks of \p RTChecks to \p Plan.
   void attachRuntimeChecks(VPlan &Plan, GeneratedRTChecks &RTChecks,
                            bool HasBranchWeights) const;

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -4092,6 +4092,7 @@ static bool willGenerateVectors(VPlan &Plan, ElementCount VF,
       case VPDef::VPEVLBasedIVPHISC:
       case VPDef::VPPredInstPHISC:
       case VPDef::VPBranchOnMaskSC:
+      case VPDef::VPScalarIVPromotionRecipeSC:
         continue;
       case VPDef::VPReductionSC:
       case VPDef::VPActiveLaneMaskPHISC:
@@ -7523,6 +7524,14 @@ BasicBlock *EpilogueVectorizerEpilogueLoop::createVectorizedLoopSkeleton() {
   OriginalScalarPH->setName("vec.epilog.iter.check");
   VPIRBasicBlock *NewEntry = Plan.createVPIRBasicBlock(OriginalScalarPH);
   VPBasicBlock *OldEntry = Plan.getEntry();
+
+  for (VPRecipeBase &R : make_early_inc_range(*OldEntry))
+    // Move hoisted loads to split PreHeader
+    if (auto RepR = dyn_cast<VPReplicateRecipe>(&R)) {
+      RepR->removeFromParent();
+      VectorPHVPBB->appendRecipe(RepR);
+    }
+
   for (auto &R : make_early_inc_range(*OldEntry)) {
     // Skip moving VPIRInstructions (including VPIRPhis), which are unmovable by
     // defining.
@@ -7532,6 +7541,7 @@ BasicBlock *EpilogueVectorizerEpilogueLoop::createVectorizedLoopSkeleton() {
   }
 
   VPBlockUtils::reassociateBlocks(OldEntry, NewEntry);
+
   Plan.setEntry(NewEntry);
   // OldEntry is now dead and will be cleaned up when the plan gets destroyed.
 
@@ -8324,6 +8334,23 @@ void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
   }
 }
 
+void LoopVectorizationPlanner::adjustScalarIVPromotions(VPlanPtr &Plan) {
+  VPScalarIVPromotionRecipe *Recipe = nullptr;
+
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           vp_depth_first_deep(Plan->getVectorLoopRegion())))
+    for (VPRecipeBase &R : *VPBB)
+      if (auto *ScalarIV = dyn_cast<VPScalarIVPromotionRecipe>(&R)) {
+        assert(!Recipe && "Only one FFLoad is supported");
+        Recipe = ScalarIV;
+      }
+
+  if (!Recipe)
+    return;
+
+  Recipe->setVFxUF(&Plan->getVFxUF());
+}
+
 VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
     VPlanPtr Plan, VFRange &Range, LoopVersioning *LVer) {
 
@@ -8434,11 +8461,12 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
       // latter are added above for masking.
       // FIXME: Migrate code relying on the underlying instruction from VPlan0
       // to construct recipes below to not use the underlying instruction.
-      if (isa<VPCanonicalIVPHIRecipe, VPWidenCanonicalIVRecipe, VPBlendRecipe>(
-              &R) ||
+      if (isa<VPCanonicalIVPHIRecipe, VPWidenCanonicalIVRecipe, VPBlendRecipe,
+              VPScalarIVPromotionRecipe>(&R) ||
           (isa<VPInstruction>(&R) && !UnderlyingValue))
         continue;
-      assert(isa<VPInstruction>(&R) && UnderlyingValue && "unsupported recipe");
+      assert((isa<VPInstruction, VPReplicateRecipe>(&R) && UnderlyingValue &&
+              "unsupported recipe"));
 
       // TODO: Gradually replace uses of underlying instruction by analyses on
       // VPlan.
@@ -8514,6 +8542,8 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
   // Adjust the recipes for any inloop reductions.
   adjustRecipesForReductions(Plan, RecipeBuilder, Range.Start);
 
+  adjustScalarIVPromotions(Plan);
+
   // Apply mandatory transformation to handle FP maxnum/minnum reduction with
   // NaNs if possible, bail out otherwise.
   if (!VPlanTransforms::runPass(VPlanTransforms::handleMaxMinNumReductions,

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -554,6 +554,7 @@ class VPSingleDefRecipe : public VPRecipeBase, public VPValue {
     case VPRecipeBase::VPWidenPointerInductionSC:
     case VPRecipeBase::VPReductionPHISC:
     case VPRecipeBase::VPPartialReductionSC:
+    case VPRecipeBase::VPScalarIVPromotionRecipeSC:
       return true;
     case VPRecipeBase::VPBranchOnMaskSC:
     case VPRecipeBase::VPInterleaveEVLSC:
@@ -580,10 +581,12 @@ class VPSingleDefRecipe : public VPRecipeBase, public VPValue {
 
   /// Returns the underlying instruction.
   Instruction *getUnderlyingInstr() {
-    return cast<Instruction>(getUnderlyingValue());
+    return getUnderlyingValue() ? dyn_cast<Instruction>(getUnderlyingValue())
+                                : nullptr;
   }
   const Instruction *getUnderlyingInstr() const {
-    return cast<Instruction>(getUnderlyingValue());
+    return getUnderlyingValue() ? dyn_cast<Instruction>(getUnderlyingValue())
+                                : nullptr;
   }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -2312,7 +2315,8 @@ struct VPFirstOrderRecurrencePHIRecipe : public VPHeaderPHIRecipe {
 
   VPFirstOrderRecurrencePHIRecipe *clone() override {
     return new VPFirstOrderRecurrencePHIRecipe(
-        cast<PHINode>(getUnderlyingInstr()), *getOperand(0));
+        getUnderlyingInstr() ? cast<PHINode>(getUnderlyingInstr()) : nullptr,
+        *getOperand(0));
   }
 
   void execute(VPTransformState &State) override;
@@ -3475,6 +3479,54 @@ class VPExpandSCEVRecipe : public VPSingleDefRecipe {
   const SCEV *getSCEV() const { return Expr; }
 };
 
+struct LLVM_ABI_FOR_TEST VPScalarIVPromotionRecipe : public VPSingleDefRecipe {
+  VPScalarIVPromotionRecipe(std::initializer_list<VPValue *> Operands,
+                            DebugLoc DL = DebugLoc::getUnknown())
+      : VPSingleDefRecipe(VPDef::VPScalarIVPromotionRecipeSC, Operands, DL) {}
+
+  VP_CLASSOF_IMPL(VPDef::VPScalarIVPromotionRecipeSC)
+
+  bool isSingleScalar() const { return true; }
+
+  VPScalarIVPromotionRecipe *clone() override {
+    assert(getNumOperands() == 3 || getNumOperands() == 4);
+    if (getNumOperands() == 3)
+      return new VPScalarIVPromotionRecipe(
+          {getOperand(0), getOperand(1), getOperand(2)}, getDebugLoc());
+    return new VPScalarIVPromotionRecipe(
+        {getOperand(0), getOperand(1), getOperand(2), getOperand(3)},
+        getDebugLoc());
+  }
+
+  VPValue *getVFxUF() { return getOperand(3); }
+  void setVFxUF(VPValue *V) {
+    if (getNumOperands() == 3) {
+      addOperand(V);
+    } else {
+      setOperand(3, V);
+    }
+  }
+
+  void execute(VPTransformState &State) override;
+
+  InstructionCost computeCost(ElementCount VF,
+                              VPCostContext &Ctx) const override {
+    return 0;
+  }
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// Print the recipe.
+  void print(raw_ostream &O, const Twine &Indent,
+             VPSlotTracker &SlotTracker) const override;
+#endif
+
+  bool usesScalars(const VPValue *Op) const override {
+    assert(is_contained(operands(), Op) &&
+           "Op must be an operand of the recipe");
+    return true;
+  }
+};
+
 /// Canonical scalar induction phi of the vector loop. Starting at the specified
 /// start value (either 0 or the resume value when vectorizing the epilogue
 /// loop). VPWidenCanonicalIVRecipe represents the vector version of the

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -136,6 +136,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     return inferScalarType(R->getOperand(0));
   case VPInstruction::BranchOnCond:
   case VPInstruction::BranchOnCount:
+  case Instruction::Store:
     return Type::getVoidTy(Ctx);
   default:
     break;
@@ -289,9 +290,10 @@ Type *VPTypeAnalysis::inferScalarType(const VPValue *V) {
           .Case<VPReductionRecipe, VPPredInstPHIRecipe, VPWidenPHIRecipe,
                 VPScalarIVStepsRecipe, VPWidenGEPRecipe, VPVectorPointerRecipe,
                 VPVectorEndPointerRecipe, VPWidenCanonicalIVRecipe,
-                VPPartialReductionRecipe>([this](const VPRecipeBase *R) {
-            return inferScalarType(R->getOperand(0));
-          })
+                VPPartialReductionRecipe, VPScalarIVPromotionRecipe>(
+              [this](const VPRecipeBase *R) {
+                return inferScalarType(R->getOperand(0));
+              })
           // VPInstructionWithType must be handled before VPInstruction.
           .Case<VPInstructionWithType, VPWidenIntrinsicRecipe,
                 VPWidenCastRecipe>(