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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: Nicholas Guy (NickGuy-Arm)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/136680.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+7-3) 
- (modified) llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll (+22) 


``````````diff
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) }
 
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/136680