[llvm] [VPlan] Don't cost scalar IV steps if the vector body runs at most once. (PR #221974)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/221974

>From e1662a794b950d0db44866a57e47310b97fba539 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 5 Sep 2026 17:15:18 +0100
Subject: [PATCH] [VPlan] Don't cost scalar IV steps if the vector body runs at
 most once.

If the vector loop executes exactly once, each scalar IV step folds to a
constant, and there's nothing to cost.

Similarly to BranchOnCount and VPRegionBlock::cost, treat
VPScalarIVSteps as free in that case.

This fixes a slight inaccuracy introduced when moving to computing
VPScalarIVSteps costs in VPlan directly.
---
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp               | 5 +++++
 .../Transforms/LoopVectorize/AArch64/scalar-steps-cost.ll    | 2 +-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index ff212663a3fac..7d671e6cd29f1 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -3184,6 +3184,11 @@ InstructionCost VPScalarIVStepsRecipe::computeCost(ElementCount VF,
   if (vputils::onlyFirstLaneUsed(this))
     return 0;
 
+  // If the vector body executes at most once, the canonical IV is a constant
+  // and every lane's step folds away with it.
+  if (VPCostContext::executesAtMostOnce(*getParent()->getPlan(), VF))
+    return 0;
+
   // Typically the operations are:
   //   1. Add the start index to each lane value.
   //   2. Multiply the start index by the step.
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/scalar-steps-cost.ll b/llvm/test/Transforms/LoopVectorize/AArch64/scalar-steps-cost.ll
index 5c12dc366b43c..00b50c5d327b6 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/scalar-steps-cost.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/scalar-steps-cost.ll
@@ -54,7 +54,7 @@ exit:
 define i32 @scalar_steps_all_lanes_tc_eq_vf(ptr %start) {
 ; CHECK-LABEL: LV: Checking a loop in 'scalar_steps_all_lanes_tc_eq_vf'
 ; CHECK: Cost of 1 for VF 2: {{.*}} = SCALAR-STEPS {{.*}}, ir<1>, {{.*}}
-; CHECK: Cost of 1 for VF 4: {{.*}} = SCALAR-STEPS {{.*}}, ir<1>, {{.*}}
+; CHECK: Cost of 0 for VF 4: {{.*}} = SCALAR-STEPS {{.*}}, ir<1>, {{.*}}
 entry:
   br label %loop