[llvm-branch-commits] [llvm] [LV][REVEC] Correctly compute register usage (PR #213987)

[Gaëtan Bossu via llvm-branch-commits]

https://github.com/gbossu updated https://github.com/llvm/llvm-project/pull/213987

>From 40338c54ed0cea71df1ba3ea5f37f30cdd570bb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Bossu?= <gaetan.bossu at arm.com>
Date: Tue, 4 Aug 2026 15:01:21 +0000
Subject: [PATCH] [LV][REVEC] Correctly compute register usage

For REVEC, the initial types might already be vectors, so make sure the
right register class is picked.
---
 .../Transforms/Vectorize/VPlanAnalysis.cpp    | 29 ++++++++------
 .../LoopVectorize/revec-reg-usage.ll          | 39 +++++++++++++++++++
 2 files changed, 56 insertions(+), 12 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/revec-reg-usage.ll

diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index 77a33339eb5f9..969be883e353d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -189,11 +189,12 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
 
   const auto &TTICapture = TTI;
   auto GetRegUsage = [&TTICapture](Type *Ty, ElementCount VF) -> unsigned {
-    if (Ty->isTokenTy() || !VectorType::isValidElementType(Ty) ||
+    Type *EltTy = Ty->getScalarType();
+    if (Ty->isTokenTy() || !VectorType::isValidElementType(EltTy) ||
         (VF.isScalable() &&
-         !TTICapture.isElementTypeLegalForScalableVector(Ty)))
+         !TTICapture.isElementTypeLegalForScalableVector(EltTy)))
       return 0;
-    return TTICapture.getRegUsageForType(VectorType::get(Ty, VF));
+    return TTICapture.getRegUsageForType(toVectorTy(Ty, VF));
   };
 
   VPValue *CanIV = LoopRegion->getCanonicalIV();
@@ -248,6 +249,10 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
             match(VPV, m_ExtractLastPart(m_VPValue())))
           continue;
 
+        // For REVEC, the initial type might be a vector.
+        Type *InitialTy = VPV->getScalarType();
+        Type *EltTy = InitialTy->getScalarType();
+
         if (VFs[J].isScalar() ||
             isa<VPRegionValue, VPReplicateRecipe, VPDerivedIVRecipe,
                 VPCurrentIterationPHIRecipe, VPScalarIVStepsRecipe>(VPV) ||
@@ -255,7 +260,7 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
             (isa<VPReductionPHIRecipe>(VPV) &&
              (cast<VPReductionPHIRecipe>(VPV))->isInLoop())) {
           unsigned ClassID =
-              TTI.getRegisterClassForType(false, VPV->getScalarType());
+              TTI.getRegisterClassForType(InitialTy->isVectorTy(), EltTy);
           // FIXME: The target might use more than one register for the type
           // even in the scalar case.
           RegUsage[ClassID] += 1;
@@ -271,9 +276,8 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
                               << " to " << VF << " for " << *R << "\n";);
           }
 
-          Type *ScalarTy = VPV->getScalarType();
-          unsigned ClassID = TTI.getRegisterClassForType(true, ScalarTy);
-          RegUsage[ClassID] += GetRegUsage(ScalarTy, VF);
+          unsigned ClassID = TTI.getRegisterClassForType(true, EltTy);
+          RegUsage[ClassID] += GetRegUsage(InitialTy, VF);
         }
       }
 
@@ -307,12 +311,13 @@ SmallVector<VPRegisterUsage, 8> llvm::calculateRegisterUsageForPlan(
     for (auto *In : LoopInvariants) {
       // FIXME: The target might use more than one register for the type
       // even in the scalar case.
-      bool IsScalar = vputils::onlyScalarValuesUsed(In);
+      bool OnlyFirstLane = vputils::onlyScalarValuesUsed(In);
+      Type *InitialTy = In->getScalarType();
 
-      ElementCount VF = IsScalar ? ElementCount::getFixed(1) : VFs[Idx];
-      unsigned ClassID =
-          TTI.getRegisterClassForType(VF.isVector(), In->getScalarType());
-      Invariant[ClassID] += GetRegUsage(In->getScalarType(), VF);
+      ElementCount VF = OnlyFirstLane ? ElementCount::getFixed(1) : VFs[Idx];
+      unsigned ClassID = TTI.getRegisterClassForType(
+          VF.isVector() || InitialTy->isVectorTy(), InitialTy->getScalarType());
+      Invariant[ClassID] += GetRegUsage(InitialTy, VF);
     }
 
     LLVM_DEBUG({
diff --git a/llvm/test/Transforms/LoopVectorize/revec-reg-usage.ll b/llvm/test/Transforms/LoopVectorize/revec-reg-usage.ll
new file mode 100644
index 0000000000000..89f6c3969bf3c
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/revec-reg-usage.ll
@@ -0,0 +1,39 @@
+; REQUIRES: asserts
+; RUN: opt -disable-output -passes=loop-vectorize -vectorize-vector-loops \
+; RUN:     -force-vector-width=1 -debug-only=vplan < %s 2>&1 \
+; RUN:     | FileCheck %s
+; RUN: opt -disable-output -passes=loop-vectorize -vectorize-vector-loops \
+; RUN:     -force-vector-width="vscale x 1" -debug-only=vplan < %s 2>&1 \
+; RUN:     | FileCheck %s
+
+; When re-vectorising with VF = vscale x 1, the number of used registers is
+; expected to remain the same, as we are turning fixed vectors into scalable
+; ones and they belong to the same class according to the generic getRegisterClassForType.
+
+; CHECK:      LV(REG): Found max usage: 2 item
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 3 registers
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 2 registers
+; CHECK-NEXT: LV(REG): Found invariant usage: 2 item
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 1 registers
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 1 registers
+define void @register_usage(ptr noalias %dst, ptr noalias %src,
+                            ptr noalias %threshold.ptr, i64 %n) {
+entry:
+  %threshold = load <4 x i32>, ptr %threshold.ptr, align 16
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %src.gep = getelementptr inbounds <4 x i32>, ptr %src, i64 %iv
+  %dst.gep = getelementptr inbounds <4 x i32>, ptr %dst, i64 %iv
+  %v = load <4 x i32>, ptr %src.gep, align 16
+  %cmp = icmp sgt <4 x i32> %v, %threshold
+  %sel = select <4 x i1> %cmp, <4 x i32> %v, <4 x i32> %threshold
+  store <4 x i32> %sel, ptr %dst.gep, align 16
+  %iv.next = add nuw i64 %iv, 1
+  %done = icmp eq i64 %iv.next, %n
+  br i1 %done, label %exit, label %loop
+
+exit:
+  ret void
+}