[llvm] [loop-vectorize] Fix crash when using types that aren't known scale factors (PR #136680)

[Nicholas Guy via llvm-commits]

https://github.com/NickGuy-Arm created https://github.com/llvm/llvm-project/pull/136680

None

>From 230b1e84353c6208fcc50351d3dfc2a8c31c2873 Mon Sep 17 00:00:00 2001
From: Nick Guy <nicholas.guy at arm.com>
Date: Tue, 22 Apr 2025 11:14:18 +0100
Subject: [PATCH] [loop-vectorize] Fix crash when using types that aren't known
 scale factors

---
 .../Transforms/Vectorize/LoopVectorize.cpp    | 10 ++++++---
 .../AArch64/partial-reduce-chained.ll         | 22 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 08e125eca591e..7b5b8555591a2 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8826,9 +8826,13 @@ bool VPRecipeBuilder::getScaledReductions(
 
   PartialReductionChain Chain(RdxExitInstr, ExtA, ExtB, BinOp);
 
-  unsigned TargetScaleFactor =
-      PHI->getType()->getPrimitiveSizeInBits().getKnownScalarFactor(
-          A->getType()->getPrimitiveSizeInBits());
+  TypeSize PHISize = PHI->getType()->getPrimitiveSizeInBits();
+  TypeSize ASize = A->getType()->getPrimitiveSizeInBits();
+
+  if (!PHISize.hasKnownScalarFactor(ASize))
+    return false;
+
+  unsigned TargetScaleFactor = PHISize.getKnownScalarFactor(ASize);
 
   if (LoopVectorizationPlanner::getDecisionAndClampRange(
           [&](ElementCount VF) {
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
index 4e4a5c82c298a..061fa99c1ba41 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
@@ -1030,6 +1030,28 @@ for.body:                                         ; preds = %for.body.preheader,
   br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !loop !1
 }
 
+
+define void @chained_partial_reduce_not_known_factor(i32 %a, i32 %b, i32 %N) #0 {
+entry:
+  br label %for.body
+
+for.cond.cleanup:
+  %B.0.lcssa = phi i40 [ %2, %for.body ]
+  ret void
+
+for.body:
+  %B.017 = phi i40 [ 0, %entry ], [ %2, %for.body ]
+  %i.016 = phi i16 [ 0, %entry ], [ %add21, %for.body ]
+  %resize = sext i32 %a to i40
+  %resize4 = sext i32 %b to i40
+  %0 = or i40 %resize4, %resize
+  %1 = or i40 %B.017, %0
+  %2 = or i40 %1, 0
+  %add21 = add i16 %i.016, 1
+  %cmp = icmp slt i16 %i.016, 1
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+}
+
 attributes #0 = { vscale_range(1,16) }