From 130f0f526dc28ebbe23e5956857e85f7c9b754f5 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Fri, 5 Jul 2024 16:56:46 +0200
Subject: [PATCH] [LVI][CVP] Add support for vector comparisons

---
 llvm/lib/Analysis/LazyValueInfo.cpp           |  6 +---
 .../CorrelatedValuePropagation/vectors.ll     | 32 +++++++++++++++----
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index 27a25377aa86bb..877898f6daeef6 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -1784,12 +1784,8 @@ static Constant *getPredicateResult(CmpInst::Predicate Pred, Constant *C,
 
   Type *ResTy = CmpInst::makeCmpResultType(C->getType());
   if (Val.isConstantRange()) {
-    ConstantInt *CI = dyn_cast<ConstantInt>(C);
-    if (!CI)
-      return nullptr;
-
     const ConstantRange &CR = Val.getConstantRange();
-    ConstantRange RHS(CI->getValue());
+    ConstantRange RHS = C->toConstantRange();
     if (CR.icmp(Pred, RHS))
       return ConstantInt::getTrue(ResTy);
     if (CR.icmp(CmpInst::getInversePredicate(Pred), RHS))
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
index 88e995ac747ae0..6f13263fe92be8 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
@@ -1,32 +1,52 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; RUN: opt -S -passes=correlated-propagation < %s | FileCheck %s
 
-; TODO: Add support for this.
 define <2 x i1> @cmp1(<2 x i8> %a) {
 ; CHECK-LABEL: define <2 x i1> @cmp1(
 ; CHECK-SAME: <2 x i8> [[A:%.*]]) {
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw <2 x i8> [[A]], <i8 1, i8 1>
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i8> [[ADD]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+; CHECK-NEXT:    ret <2 x i1> <i1 true, i1 true>
 ;
   %add = add nuw <2 x i8> %a, splat (i8 1)
   %cmp = icmp ne <2 x i8> %add, zeroinitializer
   ret <2 x i1> %cmp
 }
 
-; TODO: Add support for this.
 define <2 x i1> @cmp2(<2 x i8> %a) {
 ; CHECK-LABEL: define <2 x i1> @cmp2(
 ; CHECK-SAME: <2 x i8> [[A:%.*]]) {
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw <2 x i8> [[A]], <i8 5, i8 5>
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt <2 x i8> [[ADD]], <i8 2, i8 2>
-; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+; CHECK-NEXT:    ret <2 x i1> <i1 true, i1 true>
 ;
   %add = add nuw <2 x i8> %a, splat (i8 5)
   %cmp = icmp ugt <2 x i8> %add, splat (i8 2)
   ret <2 x i1> %cmp
 }
 
+define <2 x i1> @cmp_nonsplat(<2 x i8> %a) {
+; CHECK-LABEL: define <2 x i1> @cmp_nonsplat(
+; CHECK-SAME: <2 x i8> [[A:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw <2 x i8> [[A]], <i8 4, i8 5>
+; CHECK-NEXT:    ret <2 x i1> <i1 true, i1 true>
+;
+  %add = add nuw <2 x i8> %a, <i8 4, i8 5>
+  %cmp = icmp ugt <2 x i8> %add, <i8 2, i8 3>
+  ret <2 x i1> %cmp
+}
+
+; Handling this would require keeping track of ranges on a per-element basis.
+define <2 x i1> @cmp_nonsplat_fail(<2 x i8> %a) {
+; CHECK-LABEL: define <2 x i1> @cmp_nonsplat_fail(
+; CHECK-SAME: <2 x i8> [[A:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw <2 x i8> [[A]], <i8 3, i8 4>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt <2 x i8> [[ADD]], <i8 2, i8 3>
+; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+;
+  %add = add nuw <2 x i8> %a, <i8 3, i8 4>
+  %cmp = icmp ugt <2 x i8> %add, <i8 2, i8 3>
+  ret <2 x i1> %cmp
+}
+
 define <2 x i1> @cmp_signedness(<2 x i8> %a) {
 ; CHECK-LABEL: define <2 x i1> @cmp_signedness(
 ; CHECK-SAME: <2 x i8> [[A:%.*]]) {