[llvm] [VPlan][LV] Fix invalid truncation in VPScalarIVStepsRecipe (PR #137832)

[Maryam Moghadas via llvm-commits]

https://github.com/maryammo created https://github.com/llvm/llvm-project/pull/137832

Replace CreateTrunc with CreateSExtOrTrunc in VPScalarIVStepsRecipe to safely handle type conversion. This prevents assertion failures from invalid truncation when StartIdx0 has a smaller integer type than IntStepTy. The assertion was introduced by commit 783a846.

>From 25c647a4f52fca87c548d1e34f55ed876967557a Mon Sep 17 00:00:00 2001
From: Maryam Moghadas <maryammo at ca.ibm.com>
Date: Tue, 29 Apr 2025 15:48:35 +0000
Subject: [PATCH] [VPlan][LV] Fix invalid truncation in VPScalarIVStepsRecipe

Replace CreateTrunc with CreateSExtOrTrunc in VPScalarIVStepsRecipe to safely
handle type conversion. This prevents assertion failures from invalid truncation
when StartIdx0 has a smaller integer type than IntStepTy. The assertion was
introduced by commit 783a846.
---
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp |  2 +-
 .../PowerPC/vplan-scalarivsext-crash.ll       | 31 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/PowerPC/vplan-scalarivsext-crash.ll

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 75d056026025a..6ca43eec63dd4 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -2134,7 +2134,7 @@ void VPScalarIVStepsRecipe::execute(VPTransformState &State) {
           Builder.CreateMul(StartIdx0, ConstantInt::get(StartIdx0->getType(),
                                                         getUnrollPart(*this)));
     }
-    StartIdx0 = Builder.CreateTrunc(StartIdx0, IntStepTy);
+    StartIdx0 = Builder.CreateSExtOrTrunc(StartIdx0, IntStepTy);
   }
 
   if (!FirstLaneOnly && State.VF.isScalable()) {
diff --git a/llvm/test/Transforms/LoopVectorize/PowerPC/vplan-scalarivsext-crash.ll b/llvm/test/Transforms/LoopVectorize/PowerPC/vplan-scalarivsext-crash.ll
new file mode 100644
index 0000000000000..66b54a29b913c
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/PowerPC/vplan-scalarivsext-crash.ll
@@ -0,0 +1,31 @@
+; RUN: opt -passes=loop-vectorize -S < %s > /dev/null
+; REQUIRES: asserts
+
+target datalayout = "E-m:a-p:32:32-Fi32-i64:64-n32"
+target triple = "powerpc-ibm-aix7.2.0.0"
+
+define void @__power_mod_NMOD_power_init(ptr %a, ptr %b, i32 %n) {
+entry:
+  br label %loop_entry
+
+loop_exit:                                        ; preds = %loop_header
+  br label %loop_entry
+
+loop_entry:                                       ; preds = %loop_exit, %entry
+  %sum.0 = phi double [ 0.000000e+00, %entry ], [ %sum.1, %loop_exit ]
+  %x = load double, ptr %a, align 8
+  br label %loop_header
+
+loop_header:                                      ; preds = %loop_body, %loop_entry
+  %sum.1 = phi double [ %sum.0, %loop_entry ], [ %sum.next, %loop_body ]
+  %i = phi i32 [ 0, %loop_entry ], [ %i.next, %loop_body ]
+  %cond = icmp sgt i32 %i, %n
+  br i1 %cond, label %loop_exit, label %loop_body
+
+loop_body:                                        ; preds = %loop_header
+  store double %sum.1, ptr %b, align 8
+  %sum.next = fadd reassoc double %sum.1, %x
+  %i.next = add i32 %i, 1
+  br label %loop_header
+}
+