[InstCombine] Improve foldAndOrOfICmpsUsingRanges

[dtcxzyw]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/67327.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp (+40-14)
• (modified) llvm/test/Transforms/InstCombine/icmp-range.ll (+2-5)

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

[nikic]

This looks like something that should be handled in foldLogOpOfMaskedICmps()? Why is this in foldAndOrOfICmpsUsingRanges()?

[dtcxzyw]

This looks like something that should be handled in foldLogOpOfMaskedICmps()? Why is this in foldAndOrOfICmpsUsingRanges()?

Because the following equivalent form is handled in foldAndOrOfICmpsUsingRanges:

define i1 @isFloat2(i64 %0) {
  %a = and i64 %0, -281474976710656
  %cmp1 = icmp ne i64 %a, 0
  %cmp2 = icmp ne i64 %a, 281474976710656
  %res = and i1 %cmp1, %cmp2
  ret i1 %res
}

godbolt: https://godbolt.org/z/j7YYhGPzr

[dtcxzyw]

Ping.

[goldsteinn]

nit: assert msg

[goldsteinn]

1. please add more robust tests. Think you need some negative ones. Maybe:

• B != D
• some of the explicit pattern (with both sides eq/ne)
• some non eq/ne patterns with a pred that can't be matched in mask pattern (maybe some signed preds, non-mask + ugt, etc...)
• some vec tests with undef elements.