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

Summary:
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.

Test Plan: 

Reviewers: 

Subscribers: 

Tasks: 

Tags: 


Differential Revision: https://phabricator.intern.facebook.com/D60250842

Author: nikic

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()));

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:

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: