[ValueTracking] Add support for non-splat vecs in computeConstantRange (#72365)

Related patch: https://github.com/llvm/llvm-project/pull/68331
This missed optimization is discovered with the help of
https://github.com/AliveToolkit/alive2/pull/962.

4 files changed, 13 insertions, 12 deletions

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index dc3c90b..5f5d7e07 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -8871,9 +8871,17 @@ ConstantRange llvm::computeConstantRange(const Value *V, bool ForSigned,
   const APInt *C;
   if (match(V, m_APInt(C)))
     return ConstantRange(*C);
+  unsigned BitWidth = V->getType()->getScalarSizeInBits();
+
+  if (auto *VC = dyn_cast<ConstantDataVector>(V)) {
+    ConstantRange CR = ConstantRange::getEmpty(BitWidth);
+    for (unsigned ElemIdx = 0, NElem = VC->getNumElements(); ElemIdx < NElem;
+         ++ElemIdx)
+      CR = CR.unionWith(VC->getElementAsAPInt(ElemIdx));
+    return CR;
+  }
 
   InstrInfoQuery IIQ(UseInstrInfo);
-  unsigned BitWidth = V->getType()->getScalarSizeInBits();
   ConstantRange CR = ConstantRange::getFull(BitWidth);
   if (auto *BO = dyn_cast<BinaryOperator>(V)) {
     APInt Lower = APInt(BitWidth, 0);
diff --git a/llvm/test/Transforms/InstCombine/add.ll b/llvm/test/Transforms/InstCombine/add.ll
index 1079f4f..2ae0181 100644
--- a/llvm/test/Transforms/InstCombine/add.ll
+++ b/llvm/test/Transforms/InstCombine/add.ll
@@ -440,11 +440,9 @@ define <2 x i8> @test18vec_nsw(<2 x i8> %A) {
   ret <2 x i8> %C
 }
 
-; TODO: fix ValueTracking overflow check for non-splat vector, could be attached nsw
-; this shouldn't overflow.
 define <2 x i8> @test18vec_nsw_false(<2 x i8> %A) {
 ; CHECK-LABEL: @test18vec_nsw_false(
-; CHECK-NEXT:    [[C:%.*]] = sub <2 x i8> <i8 -125, i8 -126>, [[A:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = sub nsw <2 x i8> <i8 -125, i8 -126>, [[A:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[C]]
 ;
   %B = xor <2 x i8> %A, <i8 -1, i8 -1>
diff --git a/llvm/test/Transforms/InstCombine/addsub-constant-folding.ll b/llvm/test/Transforms/InstCombine/addsub-constant-folding.ll
index 83a45e697e5..6a6bf74 100644
--- a/llvm/test/Transforms/InstCombine/addsub-constant-folding.ll
+++ b/llvm/test/Transforms/InstCombine/addsub-constant-folding.ll
@@ -203,11 +203,9 @@ define <2 x i8> @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov1(<2 x i8> %arg
   ret <2 x i8> %t1
 }
 
-
-; TODO: We can add nsw on sub, current Value Tracking use [max element,min element] constant range, to check overflow for vector?
 define <2 x i8> @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov2(<2 x i8> %arg) {
 ; CHECK-LABEL: @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov2(
-; CHECK-NEXT:    [[T1:%.*]] = sub <2 x i8> <i8 -126, i8 -128>, [[ARG:%.*]]
+; CHECK-NEXT:    [[T1:%.*]] = sub nsw <2 x i8> <i8 -126, i8 -128>, [[ARG:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[T1]]
 ;
   %t0 = add nsw <2 x i8> %arg, <i8 1, i8 2>
@@ -215,10 +213,9 @@ define <2 x i8> @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov2(<2 x i8> %arg
   ret <2 x i8> %t1
 }
 
-; TODO: We can add nsw on sub, curret Value Tracking can't decide this is not overflowed?
 define <2 x i8> @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov3(<2 x i8> %arg) {
 ; CHECK-LABEL: @non_splat_vec_add_nsw_const_const_sub_nsw_not_ov3(
-; CHECK-NEXT:    [[T1:%.*]] = sub <2 x i8> <i8 -120, i8 -127>, [[ARG:%.*]]
+; CHECK-NEXT:    [[T1:%.*]] = sub nsw <2 x i8> <i8 -120, i8 -127>, [[ARG:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[T1]]
 ;
   %t0 = add nsw <2 x i8> %arg, <i8 0, i8 1>
diff --git a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
index 612c5bf..c1bb694 100644
--- a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
+++ b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
@@ -1052,11 +1052,9 @@ define <2 x i8> @test_vector_usub_add_nuw_no_ov(<2 x i8> %a) {
   ret <2 x i8> %r
 }
 
-; Can be optimized if the usub.sat RHS constant range handles non-splat vectors.
 define <2 x i8> @test_vector_usub_add_nuw_no_ov_nonsplat1(<2 x i8> %a) {
 ; CHECK-LABEL: @test_vector_usub_add_nuw_no_ov_nonsplat1(
-; CHECK-NEXT:    [[B:%.*]] = add nuw <2 x i8> [[A:%.*]], <i8 10, i8 10>
-; CHECK-NEXT:    [[R:%.*]] = call <2 x i8> @llvm.usub.sat.v2i8(<2 x i8> [[B]], <2 x i8> <i8 10, i8 9>)
+; CHECK-NEXT:    [[R:%.*]] = add <2 x i8> [[A:%.*]], <i8 0, i8 1>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %b = add nuw <2 x i8> %a, <i8 10, i8 10>