Address review comments

Author: david-arm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -10165,6 +10165,12 @@ static void checkMixedPrecision(Loop *L, OptimizationRemarkEmitter *ORE) {
   }
 }
 
+/// For loops with uncountable early exits, find the cost of doing work when
+/// exiting the loop early, such as calculating the final exit values of
+/// variables used outside the loop.
+/// TODO: This is currently overly pessimistic because the loop may not take
+/// the early exit, but better to keep this conservative for now. In future,
+/// it might be possible to relax this by using branch probabilities.
 static InstructionCost calculateEarlyExitCost(LoopVectorizationCostModel &CM,
                                               VPlan &Plan, ElementCount VF) {
   InstructionCost Cost = 0;
@@ -10173,37 +10179,44 @@ static InstructionCost calculateEarlyExitCost(LoopVectorizationCostModel &CM,
   LLVM_DEBUG(
       dbgs() << "Calculating cost of work in vector early exit block:\n");
   for (auto *ExitVPBB : Plan.getExitBlocks()) {
-    for (auto *PredVPBB : ExitVPBB->getPredecessors())
+    for (auto *PredVPBB : ExitVPBB->getPredecessors()) {
+      // If the predecessor is not the middle.block, then it must be the
+      // vector.early.exit block, which may contain work to calculate the exit
+      // values of variables used outside the loop.
       if (PredVPBB != Plan.getMiddleBlock())
         for (auto &R : *(cast<VPBasicBlock>(PredVPBB)))
           Cost += R.cost(VF, CostCtx);
+    }
   }
   return Cost;
 }
 
+/// This function determines whether or not it's still profitable to vectorize
+/// the loop given the extra work we have to do outside of the loop:
+///  1. Perform the runtime checks before entering the loop to ensure it's safe
+///     to vectorize.
+///  2. In the case of loops with uncountable early exits, we may have to do
+///     extra work when exiting the loop early, such as calculating the final
+///     exit values of variables used outside the loop.
 static bool isOutsideLoopWorkProfitable(GeneratedRTChecks &Checks,
-                                        VectorizationFactor &VF, Loop *L,
-                                        const TargetTransformInfo &TTI,
+                                        VectorizationFactor &VF,
+                                        LoopVectorizationCostModel &CM,
                                         PredicatedScalarEvolution &PSE,
-                                        ScalarEpilogueLowering SEL,
-                                        std::optional<unsigned> VScale,
-                                        InstructionCost EarlyExitCost) {
-  InstructionCost CheckCost = Checks.getCost();
-  if (!CheckCost.isValid() && !EarlyExitCost.isValid())
+                                        VPlan &Plan,
+                                        ScalarEpilogueLowering SEL) {
+  InstructionCost TotalCost = Checks.getCost();
+  if (!TotalCost.isValid())
     return false;
 
-  InstructionCost TotalCost = 0;
-  if (CheckCost.isValid())
-    TotalCost += CheckCost;
-
   // Add on the cost of work required in the vector early exit block, if one
   // exists.
-  if (EarlyExitCost.isValid())
-    TotalCost += EarlyExitCost;
+  if (CM.Legal->hasUncountableEarlyExit())
+    TotalCost += calculateEarlyExitCost(CM, Plan, VF.Width);
 
   // When interleaving only scalar and vector cost will be equal, which in turn
   // would lead to a divide by 0. Fall back to hard threshold.
   if (VF.Width.isScalar()) {
+    // TODO: Should we rename VectorizeMemoryCheckThreshold?
     if (TotalCost > VectorizeMemoryCheckThreshold) {
       LLVM_DEBUG(
           dbgs()
@@ -10229,7 +10242,9 @@ static bool isOutsideLoopWorkProfitable(GeneratedRTChecks &Checks,
   //  The total cost of the vector loop is
   //    RtC + VecC * (TC / VF) + EpiC
   //  where
-  //  * RtC is the cost of the generated runtime checks
+  //  * RtC is the cost of the generated runtime checks plus the cost of
+  //    performing any additional work in the vector.early.exit block for loops
+  //    with uncountable early exits.
   //  * VecC is the cost of a single vector iteration.
   //  * TC is the actual trip count of the loop
   //  * VF is the vectorization factor
@@ -10246,7 +10261,7 @@ static bool isOutsideLoopWorkProfitable(GeneratedRTChecks &Checks,
   // For now we assume the epilogue cost EpiC = 0 for simplicity. Note that
   // the computations are performed on doubles, not integers and the result
   // is rounded up, hence we get an upper estimate of the TC.
-  unsigned IntVF = getEstimatedRuntimeVF(VF.Width, VScale);
+  unsigned IntVF = getEstimatedRuntimeVF(VF.Width, CM.getVScaleForTuning());
   uint64_t RtC = *TotalCost.getValue();
   uint64_t Div = ScalarC * IntVF - *VF.Cost.getValue();
   uint64_t MinTC1 = Div == 0 ? 0 : divideCeil(RtC * IntVF, Div);
@@ -10274,7 +10289,7 @@ static bool isOutsideLoopWorkProfitable(GeneratedRTChecks &Checks,
 
   // Skip vectorization if the expected trip count is less than the minimum
   // required trip count.
-  if (auto ExpectedTC = getSmallBestKnownTC(PSE, L)) {
+  if (auto ExpectedTC = getSmallBestKnownTC(PSE, CM.TheLoop)) {
     if (ElementCount::isKnownLT(ElementCount::getFixed(*ExpectedTC),
                                 VF.MinProfitableTripCount)) {
       LLVM_DEBUG(dbgs() << "LV: Vectorization is not beneficial: expected "
@@ -10671,17 +10686,12 @@ bool LoopVectorizePass::processLoop(Loop *L) {
     if (VF.Width.isVector() || SelectedIC > 1)
       Checks.create(L, *LVL.getLAI(), PSE.getPredicate(), VF.Width, SelectedIC);
 
-    InstructionCost EarlyExitCost = InstructionCost::getInvalid();
-    if (VF.Width.isVector() && LVL.hasUncountableEarlyExit())
-      EarlyExitCost =
-          calculateEarlyExitCost(CM, LVP.getPlanFor(VF.Width), VF.Width);
-
     // Check if it is profitable to vectorize with runtime checks.
     bool ForceVectorization =
         Hints.getForce() == LoopVectorizeHints::FK_Enabled;
     if (!ForceVectorization &&
-        !isOutsideLoopWorkProfitable(Checks, VF, L, *TTI, PSE, SEL,
-                                     CM.getVScaleForTuning(), EarlyExitCost)) {
+        !isOutsideLoopWorkProfitable(Checks, VF, CM, PSE,
+                                     LVP.getPlanFor(VF.Width), SEL)) {
       ORE->emit([&]() {
         return OptimizationRemarkAnalysisAliasing(
                    DEBUG_TYPE, "CantReorderMemOps", L->getStartLoc(),

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -742,9 +742,9 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
   case VPInstruction::ExtractFirstActive: {
     // Calculate the cost of determining the lane index.
     auto *PredTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(1)), VF);
-    IntrinsicCostAttributes Attrs(
-        Intrinsic::experimental_cttz_elts, Type::getInt64Ty(Ctx.LLVMCtx),
-        {PoisonValue::get(PredTy), ConstantInt::getTrue(Ctx.LLVMCtx)});
+    IntrinsicCostAttributes Attrs(Intrinsic::experimental_cttz_elts,
+                                  Type::getInt64Ty(Ctx.LLVMCtx),
+                                  {PredTy, Type::getInt1Ty(Ctx.LLVMCtx)});
     InstructionCost Cost = Ctx.TTI.getIntrinsicInstrCost(Attrs, Ctx.CostKind);
     // Add on the cost of extracting the element.
     auto *VecTy = toVectorTy(Ctx.Types.inferScalarType(getOperand(0)), VF);

llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll

@@ -7,8 +7,8 @@ target triple = "aarch64-unknown-linux-gnu"
 
 declare void @init_mem(ptr, i64);
 
-define i64 @same_exit_block_pre_inc_use1() #1 {
-; CHECK-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1'
+define i64 @same_exit_block_pre_inc_use1_sve() #1 {
+; CHECK-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1_sve'
 ; CHECK: LV: Selecting VF: vscale x 16
 ; CHECK: Calculating cost of work in vector early exit block:
 ; CHECK-NEXT: Cost of 6 for VF vscale x 16: EMIT vp<{{.*}}> = extract-first-active
@@ -44,4 +44,43 @@ loop.end:
   ret i64 %retval
 }
 
+define i64 @same_exit_block_pre_inc_use1_nosve() {
+; CHECK-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1_nosve'
+; CHECK: LV: Selecting VF: 16
+; CHECK: Calculating cost of work in vector early exit block:
+; CHECK-NEXT: Cost of 50 for VF 16: EMIT vp<{{.*}}> = extract-first-active
+; CHECK-NEXT: Cost of 50 for VF 16: EMIT vp<{{.*}}> = extract-first-active
+; CHECK: LV: Minimum required TC for runtime checks to be profitable:176
+; CHECK-NEXT: LV: Vectorization is not beneficial: expected trip count < minimum profitable VF (64 < 176)
+; CHECK-NEXT: LV: Too many memory checks needed.
+entry:
+  %p1 = alloca [1024 x i8]
+  %p2 = alloca [1024 x i8]
+  call void @init_mem(ptr %p1, i64 1024)
+  call void @init_mem(ptr %p2, i64 1024)
+  br label %loop
+
+loop:
+  %index = phi i64 [ %index.next, %loop.inc ], [ 3, %entry ]
+  %index2 = phi i64 [ %index2.next, %loop.inc ], [ 15, %entry ]
+  %arrayidx = getelementptr inbounds i8, ptr %p1, i64 %index
+  %ld1 = load i8, ptr %arrayidx, align 1
+  %arrayidx1 = getelementptr inbounds i8, ptr %p2, i64 %index
+  %ld2 = load i8, ptr %arrayidx1, align 1
+  %cmp3 = icmp eq i8 %ld1, %ld2
+  br i1 %cmp3, label %loop.inc, label %loop.end
+
+loop.inc:
+  %index.next = add i64 %index, 1
+  %index2.next = add i64 %index2, 2
+  %exitcond = icmp ne i64 %index.next, 67
+  br i1 %exitcond, label %loop, label %loop.end
+
+loop.end:
+  %val1 = phi i64 [ %index, %loop ], [ 67, %loop.inc ]
+  %val2 = phi i64 [ %index2, %loop ], [ 98, %loop.inc ]
+  %retval = add i64 %val1, %val2
+  ret i64 %retval
+}
+
 attributes #1 = { "target-features"="+sve" vscale_range(1,16) }