[llvm] [ValueLattice] Support constant vectors in mergeIn() (PR #99466)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/99466

This is a followup to vector support in LVI/CVP/SCCP. In mergeIn(), if one of the operands is a vector of integer constant, we should try to convert it into a constant range, in case that allows performing a range union to something better than overdefined.

>From ded96b2f13128a9e434133892bf00344e53b6656 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 18 Jul 2024 09:35:37 +0200
Subject: [PATCH] [ValueLattice] Support constant vectors in mergeIn()

This is a followup to vector support in LVI/CVP/SCCP. In mergeIn(),
if one of the operands is a vector of integer constant, we should
try to convert it into a constant range, in case that allows
performing a range union to something better than overdefined.
---
 llvm/include/llvm/Analysis/ValueLattice.h     | 34 ++++++++++++-------
 .../CorrelatedValuePropagation/vectors.ll     |  4 +--
 llvm/test/Transforms/SCCP/phis.ll             |  4 +--
 3 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ValueLattice.h b/llvm/include/llvm/Analysis/ValueLattice.h
index b81eb5f60ab7e..fa56d838a3859 100644
--- a/llvm/include/llvm/Analysis/ValueLattice.h
+++ b/llvm/include/llvm/Analysis/ValueLattice.h
@@ -281,18 +281,21 @@ class ValueLatticeElement {
     return std::nullopt;
   }
 
-  ConstantRange asConstantRange(Type *Ty, bool UndefAllowed = false) const {
-    assert(Ty->isIntOrIntVectorTy() && "Must be integer type");
+  ConstantRange asConstantRange(unsigned BW, bool UndefAllowed = false) const {
     if (isConstantRange(UndefAllowed))
       return getConstantRange();
     if (isConstant())
       return getConstant()->toConstantRange();
-    unsigned BW = Ty->getScalarSizeInBits();
     if (isUnknown())
       return ConstantRange::getEmpty(BW);
     return ConstantRange::getFull(BW);
   }
 
+  ConstantRange asConstantRange(Type *Ty, bool UndefAllowed = false) const {
+    assert(Ty->isIntOrIntVectorTy() && "Must be integer type");
+    return asConstantRange(Ty->getScalarSizeInBits(), UndefAllowed);
+  }
+
   bool markOverdefined() {
     if (isOverdefined())
       return false;
@@ -384,7 +387,9 @@ class ValueLatticeElement {
       return true;
     }
 
-    assert(isUnknown() || isUndef());
+    assert(isUnknown() || isUndef() || isConstant());
+    assert((!isConstant() || NewR.contains(getConstant()->toConstantRange())) &&
+           "Constant must be subset of new range");
 
     NumRangeExtensions = 0;
     Tag = NewTag;
@@ -426,6 +431,16 @@ class ValueLatticeElement {
         return false;
       if (RHS.isUndef())
         return false;
+      // If the constant is a vector of integers, try to treat it as a range.
+      if (getConstant()->getType()->isVectorTy() &&
+          getConstant()->getType()->getScalarType()->isIntegerTy()) {
+        ConstantRange L = getConstant()->toConstantRange();
+        ConstantRange NewR = L.unionWith(
+            RHS.asConstantRange(L.getBitWidth(), /*UndefAllowed=*/true));
+        return markConstantRange(
+            std::move(NewR),
+            Opts.setMayIncludeUndef(RHS.isConstantRangeIncludingUndef()));
+      }
       markOverdefined();
       return true;
     }
@@ -444,14 +459,9 @@ class ValueLatticeElement {
       return OldTag != Tag;
     }
 
-    if (!RHS.isConstantRange()) {
-      // We can get here if we've encountered a constantexpr of integer type
-      // and merge it with a constantrange.
-      markOverdefined();
-      return true;
-    }
-
-    ConstantRange NewR = getConstantRange().unionWith(RHS.getConstantRange());
+    const ConstantRange &L = getConstantRange();
+    ConstantRange NewR = L.unionWith(
+        RHS.asConstantRange(L.getBitWidth(), /*UndefAllowed=*/true));
     return markConstantRange(
         std::move(NewR),
         Opts.setMayIncludeUndef(RHS.isConstantRangeIncludingUndef()));
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
index 43e680cd25cdb..6254b54d42554 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
@@ -286,7 +286,7 @@ define <2 x i16> @phi_merge1(i1 %c, <2 x i8> %a) {
 ; CHECK-NEXT:    br label %[[JOIN]]
 ; CHECK:       [[JOIN]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi <2 x i16> [ [[ZEXT]], %[[ENTRY]] ], [ <i16 1, i16 2>, %[[IF]] ]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i16> [[PHI]], <i16 2, i16 3>
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <2 x i16> [[PHI]], <i16 2, i16 3>
 ; CHECK-NEXT:    ret <2 x i16> [[ADD]]
 ;
 entry:
@@ -312,7 +312,7 @@ define <2 x i16> @phi_merge2(i1 %c, <2 x i8> %a) {
 ; CHECK-NEXT:    br label %[[JOIN]]
 ; CHECK:       [[JOIN]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi <2 x i16> [ <i16 1, i16 2>, %[[ENTRY]] ], [ [[ZEXT]], %[[IF]] ]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i16> [[PHI]], <i16 2, i16 3>
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <2 x i16> [[PHI]], <i16 2, i16 3>
 ; CHECK-NEXT:    ret <2 x i16> [[ADD]]
 ;
 entry:
diff --git a/llvm/test/Transforms/SCCP/phis.ll b/llvm/test/Transforms/SCCP/phis.ll
index 83daae0a7c0c8..9264a6eaefb85 100644
--- a/llvm/test/Transforms/SCCP/phis.ll
+++ b/llvm/test/Transforms/SCCP/phis.ll
@@ -109,7 +109,7 @@ define <2 x i16> @phi_vector_merge1(i1 %c, <2 x i8> %a) {
 ; CHECK-NEXT:    br label %[[JOIN]]
 ; CHECK:       [[JOIN]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi <2 x i16> [ [[ZEXT]], %[[ENTRY]] ], [ <i16 1, i16 2>, %[[IF]] ]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i16> [[PHI]], <i16 2, i16 3>
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <2 x i16> [[PHI]], <i16 2, i16 3>
 ; CHECK-NEXT:    ret <2 x i16> [[ADD]]
 ;
 entry:
@@ -135,7 +135,7 @@ define <2 x i16> @phi_vector_merge2(i1 %c, <2 x i8> %a) {
 ; CHECK-NEXT:    br label %[[JOIN]]
 ; CHECK:       [[JOIN]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi <2 x i16> [ <i16 1, i16 2>, %[[ENTRY]] ], [ [[ZEXT]], %[[IF]] ]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i16> [[PHI]], <i16 2, i16 3>
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <2 x i16> [[PHI]], <i16 2, i16 3>
 ; CHECK-NEXT:    ret <2 x i16> [[ADD]]
 ;
 entry: