[llvm] [loop-vectorize] Fix crash when using types that aren't known scale factors (PR #136680)
Nicholas Guy via llvm-commits
llvm-commits at lists.llvm.org
Tue Apr 22 06:46:43 PDT 2025
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/3] [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/3] 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) }
>From f33ed38fa0d6a126e9f49db170c8b1bce11fa9cc Mon Sep 17 00:00:00 2001
From: Nick Guy <nicholas.guy at arm.com>
Date: Tue, 22 Apr 2025 14:40:44 +0100
Subject: [PATCH 3/3] Update and move test
---
.../AArch64/partial-reduce-chained.ll | 20 -------
.../AArch64/partial-reduce-no-dotprod.ll | 56 +++++++++++++++++++
2 files changed, 56 insertions(+), 20 deletions(-)
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
index 197929337a09e..d58056a964f1b 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
@@ -1031,26 +1031,6 @@ for.body: ; preds = %for.body.preheader,
}
-define i40 @chained_partial_reduce_not_known_factor(i32 %a, i32 %b, i32 %N) {
-entry:
- br label %for.body
-
-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 %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) }
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-no-dotprod.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-no-dotprod.ll
index 3561f52df9490..ef82154dfce66 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-no-dotprod.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-no-dotprod.ll
@@ -59,3 +59,59 @@ for.body: ; preds = %for.body, %entry
for.exit: ; preds = %for.body
ret i32 %add
}
+
+; Test to ensure that we don't crash when evaluating an extend from a type
+; that is not a factor of the target type.
+define i40 @partial_reduce_not_known_factor(i32 %a, i32 %b, i16 %N) {
+; CHECK-LABEL: define i40 @partial_reduce_not_known_factor(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]], i16 [[N:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: [[SMAX:%.*]] = call i16 @llvm.smax.i16(i16 [[N]], i16 0)
+; CHECK-NEXT: [[TMP0:%.*]] = zext nneg i16 [[SMAX]] to i32
+; CHECK-NEXT: [[TMP1:%.*]] = add nuw nsw i32 [[TMP0]], 1
+; CHECK-NEXT: [[MIN_ITERS_CHECK:%.*]] = icmp ult i32 [[TMP1]], 4
+; CHECK-NEXT: br i1 [[MIN_ITERS_CHECK]], label %[[SCALAR_PH:.*]], label %[[VECTOR_PH:.*]]
+; CHECK: [[VECTOR_PH]]:
+; CHECK-NEXT: [[N_MOD_VF:%.*]] = urem i32 [[TMP1]], 4
+; CHECK-NEXT: [[N_VEC:%.*]] = sub i32 [[TMP1]], [[N_MOD_VF]]
+; CHECK-NEXT: [[BROADCAST_SPLATINSERT:%.*]] = insertelement <2 x i32> poison, i32 [[B]], i64 0
+; CHECK-NEXT: [[BROADCAST_SPLAT:%.*]] = shufflevector <2 x i32> [[BROADCAST_SPLATINSERT]], <2 x i32> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT: [[BROADCAST_SPLATINSERT1:%.*]] = insertelement <2 x i32> poison, i32 [[A]], i64 0
+; CHECK-NEXT: [[BROADCAST_SPLAT2:%.*]] = shufflevector <2 x i32> [[BROADCAST_SPLATINSERT1]], <2 x i32> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT: [[TMP2:%.*]] = trunc i32 [[N_VEC]] to i16
+; CHECK-NEXT: [[TMP3:%.*]] = sext <2 x i32> [[BROADCAST_SPLAT2]] to <2 x i40>
+; CHECK-NEXT: [[TMP4:%.*]] = sext <2 x i32> [[BROADCAST_SPLAT]] to <2 x i40>
+; CHECK-NEXT: [[TMP5:%.*]] = or <2 x i40> [[TMP4]], [[TMP3]]
+; CHECK-NEXT: br label %[[VECTOR_BODY:.*]]
+; CHECK: [[VECTOR_BODY]]:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i32 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <2 x i40> [ zeroinitializer, %[[VECTOR_PH]] ], [ [[TMP6:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI3:%.*]] = phi <2 x i40> [ zeroinitializer, %[[VECTOR_PH]] ], [ [[TMP8:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP6]] = or <2 x i40> [[VEC_PHI]], [[TMP5]]
+; CHECK-NEXT: [[TMP8]] = or <2 x i40> [[VEC_PHI3]], [[TMP5]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i32 [[INDEX]], 4
+; CHECK-NEXT: [[TMP7:%.*]] = icmp eq i32 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT: br i1 [[TMP7]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
+; CHECK: [[MIDDLE_BLOCK]]:
+; CHECK-NEXT: [[BIN_RDX:%.*]] = or <2 x i40> [[TMP8]], [[TMP6]]
+; CHECK-NEXT: [[TMP9:%.*]] = call i40 @llvm.vector.reduce.or.v2i40(<2 x i40> [[BIN_RDX]])
+; CHECK-NEXT: [[CMP_N:%.*]] = icmp eq i32 [[TMP1]], [[N_VEC]]
+; CHECK-NEXT: br i1 [[CMP_N]], label %[[EXIT:.*]], label %[[SCALAR_PH]]
+entry:
+ br label %for.body
+
+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 %red, %0
+ %iv.next = add i16 %iv, 1
+ %cmp = icmp slt i16 %iv, %N
+ br i1 %cmp, label %for.body, label %exit
+
+exit:
+ %result.lcssa = phi i40 [ %1, %for.body ]
+ ret i40 %result.lcssa
+}
More information about the llvm-commits
mailing list