[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 updated https://github.com/llvm/llvm-project/pull/136680

>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 1/2] [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) }
 
 

>From f213f982cbea07a4ba3f0bf59eed92b3b714c1f1 Mon Sep 17 00:00:00 2001
From: Nick Guy <nicholas.guy at arm.com>
Date: Tue, 22 Apr 2025 11:37:35 +0100
Subject: [PATCH 2/2] Update test

---
 .../AArch64/partial-reduce-chained.ll         | 23 +++++++++----------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
index 061fa99c1ba41..197929337a09e 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
@@ -1031,25 +1031,24 @@ for.body:                                         ; preds = %for.body.preheader,
 }
 
 
-define void @chained_partial_reduce_not_known_factor(i32 %a, i32 %b, i32 %N) #0 {
+define i40 @chained_partial_reduce_not_known_factor(i32 %a, i32 %b, i32 %N) {
 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 ]
+  %red = phi i40 [ 0, %entry ], [ %1, %for.body ]
+  %iv = phi i16 [ 0, %entry ], [ %iv.next, %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
+  %1 = or i40 %red, %0
+  %iv.next = add i16 %iv, 1
+  %cmp = icmp slt i16 %iv, 1
+  br i1 %cmp, label %for.body, label %exit
+
+exit:
+  %result.lcssa = phi i40 [ %1, %for.body ]
+  ret i40 %result.lcssa
 }
 
 attributes #0 = { vscale_range(1,16) }