[llvm] [LV][IRBuilder] Allow implicit truncation of step vector (PR #173229)

[Nikita Popov via llvm-commits]

https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/173229

>From c5f32bd3e83980fa3d3bc4ae5544d9ce6beb6dd2 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 22 Dec 2025 10:26:48 +0100
Subject: [PATCH] [LV][IRBuilder] Allow implicit truncation of step vector

---
 llvm/lib/IR/IRBuilder.cpp                     |  4 +-
 .../LoopVectorize/step-vector-i1-wrapping.ll  | 56 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/step-vector-i1-wrapping.ll

diff --git a/llvm/lib/IR/IRBuilder.cpp b/llvm/lib/IR/IRBuilder.cpp
index 8e1707ac98a51..a397f6973ee05 100644
--- a/llvm/lib/IR/IRBuilder.cpp
+++ b/llvm/lib/IR/IRBuilder.cpp
@@ -155,9 +155,11 @@ Value *IRBuilderBase::CreateStepVector(Type *DstType, const Twine &Name) {
   unsigned NumEls = cast<FixedVectorType>(DstType)->getNumElements();
 
   // Create a vector of consecutive numbers from zero to VF.
+  // It's okay if the values wrap around.
   SmallVector<Constant *, 8> Indices;
   for (unsigned i = 0; i < NumEls; ++i)
-    Indices.push_back(ConstantInt::get(STy, i));
+    Indices.push_back(
+        ConstantInt::get(STy, i, /*IsSigned=*/false, /*ImplicitTrunc=*/true));
 
   // Add the consecutive indices to the vector value.
   return ConstantVector::get(Indices);
diff --git a/llvm/test/Transforms/LoopVectorize/step-vector-i1-wrapping.ll b/llvm/test/Transforms/LoopVectorize/step-vector-i1-wrapping.ll
new file mode 100644
index 0000000000000..b403b300835a8
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/step-vector-i1-wrapping.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=loop-vectorize -force-vector-width=4 < %s | FileCheck %s
+
+define void @test() {
+; CHECK-LABEL: define void @test() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[VECTOR_PH:.*]]
+; CHECK:       [[VECTOR_PH]]:
+; CHECK-NEXT:    br label %[[VECTOR_BODY:.*]]
+; CHECK:       [[VECTOR_BODY]]:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VEC_IND:%.*]] = phi <4 x i1> [ <i1 false, i1 true, i1 false, i1 true>, %[[VECTOR_PH]] ], [ [[VEC_IND_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = select <4 x i1> [[VEC_IND]], <4 x i32> splat (i32 1), <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <4 x i1> [[VEC_IND]], zeroinitializer
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[INDEX_NEXT]], -4
+; CHECK-NEXT:    br i1 [[TMP1]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       [[MIDDLE_BLOCK]]:
+; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <4 x i32> [[TMP0]], i32 3
+; CHECK-NEXT:    br label %[[SCALAR_PH:.*]]
+; CHECK:       [[SCALAR_PH]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ -3, %[[SCALAR_PH]] ], [ [[IV_INC:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV2:%.*]] = phi i1 [ false, %[[SCALAR_PH]] ], [ [[IV2_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV3:%.*]] = phi i32 [ [[VECTOR_RECUR_EXTRACT]], %[[SCALAR_PH]] ], [ [[IV3_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV3_NEXT]] = select i1 [[IV2]], i32 1, i32 0
+; CHECK-NEXT:    [[IV2_NEXT]] = xor i1 [[IV2]], true
+; CHECK-NEXT:    [[IV_INC]] = add i64 [[IV]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[IV_INC]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label %[[EXIT:.*]], label %[[LOOP]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 1, %entry ], [ %iv.inc, %loop ]
+  %iv2 = phi i1 [ false, %entry ], [ %iv2.next, %loop ]
+  %iv3 = phi i32 [ 0, %entry ], [ %iv3.next, %loop ]
+  %iv3.next = select i1 %iv2, i32 1, i32 0
+  %iv2.next = xor i1 %iv2, true
+  %iv.inc = add i64 %iv, 1
+  %cmp = icmp eq i64 %iv.inc, 0
+  br i1 %cmp, label %exit, label %loop
+
+exit:
+  ret void
+}
+;.
+; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]], [[META2:![0-9]+]]}
+; CHECK: [[META1]] = !{!"llvm.loop.isvectorized", i32 1}
+; CHECK: [[META2]] = !{!"llvm.loop.unroll.runtime.disable"}
+; CHECK: [[LOOP3]] = distinct !{[[LOOP3]], [[META2]], [[META1]]}
+;.