[llvm] [VPlan] Fix register usage of VPExpressionRecipe (PR #176971)

[John Brawn via llvm-commits]

https://github.com/john-brawn-arm created https://github.com/llvm/llvm-project/pull/176971

In calculateRegisterUsageForPlan the VF is scaled when the recipe has a scaling factor. For VPExpressionRecipe this is based on the output of the recipe, but for calculating the register usage we're intrested in the largest VF used in all of the subexpressions, which means we should be using the minimum scaling factor from all of the subexpressions.

>From 8dad11794a0a51523ddc6667c9fadf529f6ed3ae Mon Sep 17 00:00:00 2001
From: John Brawn <john.brawn at arm.com>
Date: Tue, 20 Jan 2026 15:06:09 +0000
Subject: [PATCH] [VPlan] Fix register usage of VPExpressionRecipe

In calculateRegisterUsageForPlan the VF is scaled when the recipe has a scaling
factor. For VPExpressionRecipe this is based on the output of the recipe, but
for calculating the register usage we're intrested in the largest VF used in all
of the subexpressions, which means we should be using the minimum scaling factor
from all of the subexpressions.
---
 llvm/lib/Transforms/Vectorize/VPlan.h              |  2 ++
 llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp    | 14 ++++++++++----
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp     |  8 ++++++++
 .../Transforms/LoopVectorize/AArch64/reg-usage.ll  |  2 +-
 4 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 94a19beb75a8f..9ddc0a849654e 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -3215,6 +3215,8 @@ class VPExpressionRecipe : public VPSingleDefRecipe {
     return PR ? PR->getVFScaleFactor() : 1;
   }
 
+  unsigned getMinVFScaleFactor() const;
+
   /// Method for generating code, must not be called as this recipe is abstract.
   void execute(VPTransformState &State) override {
     llvm_unreachable("recipe must be removed before execute");
diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index 39e7a044dfff5..08917d72d018c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -553,10 +553,16 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
           // even in the scalar case.
           RegUsage[ClassID] += 1;
         } else {
-          // The output from scaled phis and scaled reductions actually has
-          // fewer lanes than the VF.
-          unsigned ScaleFactor =
-              vputils::getVFScaleFactor(VPV->getDefiningRecipe());
+          unsigned ScaleFactor;
+          auto *DefR = VPV->getDefiningRecipe();
+          if (auto *ER = dyn_cast<VPExpressionRecipe>(DefR))
+            // In an expression recipe we want the maximum register usage out of
+            // its subexpressions, which means the minimum scaling factor.
+            ScaleFactor = ER->getMinVFScaleFactor();
+          else
+            // The output from scaled phis and scaled reductions actually has
+            // fewer lanes than the VF.
+            ScaleFactor = vputils::getVFScaleFactor(DefR);
           ElementCount VF = VFs[J];
           if (ScaleFactor > 1) {
             VF = VFs[J].divideCoefficientBy(ScaleFactor);
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 64890ac9e9b92..fa59a2741a581 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -2949,6 +2949,14 @@ void VPExpressionRecipe::decompose() {
   ExpressionRecipes.clear();
 }
 
+unsigned VPExpressionRecipe::getMinVFScaleFactor() const {
+  unsigned Min = UINT_MAX;
+  for (auto *R : ExpressionRecipes) {
+    Min = std::min(Min, vputils::getVFScaleFactor(R));
+  }
+  return Min;
+}
+
 InstructionCost VPExpressionRecipe::computeCost(ElementCount VF,
                                                 VPCostContext &Ctx) const {
   Type *RedTy = Ctx.Types.inferScalarType(this);
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll b/llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll
index d4fb3d70c538d..45c29456ab612 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll
@@ -192,7 +192,7 @@ define i32 @dotp_unrolled(i32 %num_out, i64 %num_in, ptr %a, ptr %b) {
 ; CHECK:       LV(REG): VF = 16
 ; CHECK-NEXT:  LV(REG): Found max usage: 2 item
 ; CHECK-NEXT:  LV(REG): RegisterClass: Generic::ScalarRC, 9 registers
-; CHECK-NEXT:  LV(REG): RegisterClass: Generic::VectorRC, 6 registers
+; CHECK-NEXT:  LV(REG): RegisterClass: Generic::VectorRC, 16 registers
 ; CHECK-NEXT:  LV(REG): Found invariant usage: 1 item
 entry:
   br label %for.body