[llvm] [LVI] Generalize mask not equal conditions handling (PR #92946)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto created https://github.com/llvm/llvm-project/pull/92946

Extend `V & Mask != 0` for non-zero constants, when retrieving constraint value information from a non-equality comparison.

Proofs: https://alive2.llvm.org/ce/z/t94RAM
Motivating example: https://github.com/gcc-mirror/gcc/blob/master/gcc/testsuite/gcc.dg/tree-ssa/vrp76.c

>From 36a4f3ac1c58fe49bfbcdcebf61b0456fea39a60 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Tue, 21 May 2024 19:21:19 +0200
Subject: [PATCH] [LVI] Generalize mask not equal conditions handling

Extend `V & Mask != 0` for non-zero constants, when retrieving
constraint value information from a non-equality comparison.

Proofs: https://alive2.llvm.org/ce/z/t94RAM.
---
 llvm/lib/Analysis/LazyValueInfo.cpp           | 29 ++++++++------
 .../CorrelatedValuePropagation/icmp.ll        | 39 +++++++++++++++++--
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index 6cded828c25f4..ea43149358177 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -1183,20 +1183,27 @@ std::optional<ValueLatticeElement> LazyValueInfoImpl::getValueFromICmpCondition(
   if (match(LHS, m_And(m_Specific(Val), m_APInt(Mask))) &&
       match(RHS, m_APInt(C))) {
     // If (Val & Mask) == C then all the masked bits are known and we can
-    // compute a value range based on that.
-    if (EdgePred == ICmpInst::ICMP_EQ) {
+    // compute a value range based on that. If (Val & Mask) != C, we can return
+    // the inverse of such a range.
+    auto GetRangeFromKnownBits = [&]() -> ConstantRange {
       KnownBits Known;
       Known.Zero = ~*C & *Mask;
       Known.One = *C & *Mask;
-      return ValueLatticeElement::getRange(
-          ConstantRange::fromKnownBits(Known, /*IsSigned*/ false));
-    }
-    // If (Val & Mask) != 0 then the value must be larger than the lowest set
-    // bit of Mask.
-    if (EdgePred == ICmpInst::ICMP_NE && !Mask->isZero() && C->isZero()) {
-      return ValueLatticeElement::getRange(ConstantRange::getNonEmpty(
-          APInt::getOneBitSet(BitWidth, Mask->countr_zero()),
-          APInt::getZero(BitWidth)));
+      return ConstantRange::fromKnownBits(Known, /*IsSigned*/ false);
+    };
+
+    if (EdgePred == ICmpInst::ICMP_EQ)
+      return ValueLatticeElement::getRange(GetRangeFromKnownBits());
+
+    if (EdgePred == ICmpInst::ICMP_NE && !Mask->isZero()) {
+      // On the off-chance that (Val & Mask) != C, where C is 0, then the value
+      // must be larger than the lowest set bit of Mask.
+      if (C->isZero())
+        return ValueLatticeElement::getRange(ConstantRange::getNonEmpty(
+            APInt::getOneBitSet(BitWidth, Mask->countr_zero()),
+            APInt::getZero(BitWidth)));
+
+      return ValueLatticeElement::getRange(GetRangeFromKnownBits().inverse());
     }
   }
 
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll b/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
index b5337b9ddc248..4cfc33a94746b 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
@@ -595,7 +595,7 @@ define i1 @test_assume_cmp_with_offset_or(i64 %idx, i1 %other) {
 ; CHECK:       T:
 ; CHECK-NEXT:    ret i1 true
 ; CHECK:       F:
-; CHECK-NEXT:    ret i1 [[CMP2:%.*]]
+; CHECK-NEXT:    ret i1 [[OTHER:%.*]]
 ;
   %idx.off1 = or disjoint i64 %idx, 5
   %cmp1 = icmp ugt i64 %idx.off1, 10
@@ -1145,8 +1145,7 @@ define void @test_icmp_mask_ne_nonzero_cmp(i32 %a) {
 ; CHECK-NEXT:    call void @check1(i1 [[CMP2]])
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp ugt i32 [[A]], 2
 ; CHECK-NEXT:    call void @check1(i1 [[CMP3]])
-; CHECK-NEXT:    [[CMP4:%.*]] = icmp ult i32 [[A]], -1
-; CHECK-NEXT:    call void @check1(i1 [[CMP4]])
+; CHECK-NEXT:    call void @check1(i1 true)
 ; CHECK-NEXT:    ret void
 ; CHECK:       if.false:
 ; CHECK-NEXT:    ret void
@@ -1475,3 +1474,37 @@ entry:
   %select = select i1 %cmp1, i1 %cmp2, i1 false
   ret i1 %select
 }
+
+declare void @opaque()
+
+define void @test_icmp_ne_from_implied_range(i32 noundef %arg) {
+; CHECK-LABEL: @test_icmp_ne_from_implied_range(
+; CHECK-NEXT:    [[AND_MASK:%.*]] = and i32 [[ARG:%.*]], -8
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND_MASK]], -16
+; CHECK-NEXT:    br i1 [[CMP]], label [[END:%.*]], label [[ELSE:%.*]]
+; CHECK:       else:
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       sw.case:
+; CHECK-NEXT:    call void @opaque()
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+  %and.mask = and i32 %arg, -8
+  %cmp = icmp eq i32 %and.mask, -16
+  br i1 %cmp, label %end, label %else
+
+else:
+  switch i32 %arg, label %end [
+  i32 -16, label %sw.case
+  i32 -12, label %sw.case
+  i32 -9, label %sw.case
+  ]
+
+sw.case:
+  call void @opaque()
+  br label %end
+
+end:
+  ret void
+}