[llvm] [InstCombine] Improve `foldAndOrOfICmpsUsingRanges` (PR #67327)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/67327

This patch handles the pattern `(icmp ugt X, mask) & (icmp eq (X & ~mask), C2)` in `foldAndOrOfICmpsUsingRanges`.
It interprets the `icmp ugt X, mask` idiom into the canonical form `icmp ne (X & ~mask), 0` using `getMaskedTypeForICmpPair`.
Fixes #59555.


>From 231393577c5ce0e55bae61e7bd0855078c03464d Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 25 Sep 2023 21:03:47 +0800
Subject: [PATCH] [InstCombine] Improve `foldAndOrOfICmpsUsingRanges`

---
 .../InstCombine/InstCombineAndOrXor.cpp       | 54 ++++++++++++++-----
 .../test/Transforms/InstCombine/icmp-range.ll |  7 +--
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index c04fe827207fd0d..a092258579aa1b9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -1273,23 +1273,46 @@ Value *InstCombinerImpl::foldAndOrOfICmpsUsingRanges(ICmpInst *ICmp1,
   ICmpInst::Predicate Pred1, Pred2;
   Value *V1, *V2;
   const APInt *C1, *C2;
-  if (!match(ICmp1, m_ICmp(Pred1, m_Value(V1), m_APInt(C1))) ||
-      !match(ICmp2, m_ICmp(Pred2, m_Value(V2), m_APInt(C2))))
-    return nullptr;
-
-  // Look through add of a constant offset on V1, V2, or both operands. This
-  // allows us to interpret the V + C' < C'' range idiom into a proper range.
   const APInt *Offset1 = nullptr, *Offset2 = nullptr;
-  if (V1 != V2) {
-    Value *X;
-    if (match(V1, m_Add(m_Value(X), m_APInt(Offset1))))
-      V1 = X;
-    if (match(V2, m_Add(m_Value(X), m_APInt(Offset2))))
-      V2 = X;
+  bool Matched = false;
+
+  if (match(ICmp1, m_ICmp(Pred1, m_Value(V1), m_APInt(C1))) &&
+      match(ICmp2, m_ICmp(Pred2, m_Value(V2), m_APInt(C2)))) {
+    // Look through add of a constant offset on V1, V2, or both operands. This
+    // allows us to interpret the V + C' < C'' range idiom into a proper range.
+    if (V1 != V2) {
+      Value *X;
+      if (match(V1, m_Add(m_Value(X), m_APInt(Offset1))))
+        V1 = X;
+      if (match(V2, m_Add(m_Value(X), m_APInt(Offset2))))
+        V2 = X;
+    }
+
+    Matched = V1 == V2;
   }
 
-  if (V1 != V2)
+  if (!Matched) {
+    Value *A = nullptr, *B = nullptr, *C = nullptr, *D = nullptr, *E = nullptr;
+    // Match (icmp(A & B) ==/!= C) &/| (icmp(A & D) ==/!= E)
+    auto MaskPair =
+        getMaskedTypeForICmpPair(A, B, C, D, E, ICmp1, ICmp2, Pred1, Pred2);
+    // Match (icmp(A & B) ==/!= C1) &/| (icmp(A & B) ==/!= C2)
+    if (MaskPair && B == D && match(C, m_APIntAllowUndef(C1)) &&
+        match(E, m_APIntAllowUndef(C2)) &&
+        (match(ICmp1->getOperand(0),
+               m_CombineAnd(m_Value(V1),
+                            m_c_And(m_Specific(A), m_Specific(B)))) ||
+         match(ICmp2->getOperand(0),
+               m_CombineAnd(m_Value(V1),
+                            m_c_And(m_Specific(A), m_Specific(D)))))) {
+      V2 = V1;
+      Matched = true;
+    }
+  }
+
+  if (!Matched)
     return nullptr;
+  assert(V1 == V2);
 
   ConstantRange CR1 = ConstantRange::makeExactICmpRegion(
       IsAnd ? ICmpInst::getInversePredicate(Pred1) : Pred1, *C1);
@@ -3183,6 +3206,9 @@ Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
                               Constant::getAllOnesValue(LHS0->getType()));
   }
 
+  if (Value *V = foldAndOrOfICmpsUsingRanges(LHS, RHS, IsAnd))
+    return V;
+
   // This only handles icmp of constants: (icmp1 A, C1) | (icmp2 B, C2).
   if (!LHSC || !RHSC)
     return nullptr;
@@ -3256,7 +3282,7 @@ Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
     }
   }
 
-  return foldAndOrOfICmpsUsingRanges(LHS, RHS, IsAnd);
+  return nullptr;
 }
 
 // FIXME: We use commutative matchers (m_c_*) for some, but not all, matches
diff --git a/llvm/test/Transforms/InstCombine/icmp-range.ll b/llvm/test/Transforms/InstCombine/icmp-range.ll
index a6b629373946e30..24da5c61d741196 100644
--- a/llvm/test/Transforms/InstCombine/icmp-range.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-range.ll
@@ -1480,11 +1480,8 @@ define <2 x i1> @icmp_ne_sext_eq_otherwise_vec(<2 x i32> %a) {
 ; tests from PR59555
 define i1 @isFloat(i64 %0) {
 ; CHECK-LABEL: @isFloat(
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i64 [[TMP0:%.*]], 281474976710655
-; CHECK-NEXT:    [[TMP3:%.*]] = and i64 [[TMP0]], -281474976710656
-; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i64 [[TMP3]], 281474976710656
-; CHECK-NEXT:    [[TMP5:%.*]] = and i1 [[TMP2]], [[TMP4]]
-; CHECK-NEXT:    ret i1 [[TMP5]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i64 [[TMP0:%.*]], 562949953421311
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %2 = icmp ugt i64 %0, 281474976710655
   %3 = and i64 %0, -281474976710656