[InstCombine] Decompose an icmp into multiple ranges

[dtcxzyw]

This PR decomposes an icmp into an intersection/union of ranges. It will handle patterns that cannot be captured by foldAndOrOfICmpsUsingRanges.

Alive2: https://alive2.llvm.org/ce/z/FMK_QA
Fixes #69123.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

This PR decomposes an icmp into an intersection/union of ranges. It will handle patterns that cannot be captured by foldAndOrOfICmpsUsingRanges.

Alive2: https://alive2.llvm.org/ce/z/FMK_QA
Fixes #69123.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp (+98-1)
• (modified) llvm/test/Transforms/InstCombine/icmp-range.ll (+106)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index 3e0218d9b76d1f7..ace57fb6f597063 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -1334,6 +1334,101 @@ Value *InstCombinerImpl::foldAndOrOfICmpsUsingRanges(ICmpInst *ICmp1,
   return Builder.CreateICmp(NewPred, NewV, ConstantInt::get(Ty, NewC));
 }
 
+/// Decompose icmp into intersection or union of ranges.
+static bool decomposeICmpIntoRangeSet(SmallVectorImpl<ConstantRange> &Set,
+                                      ICmpInst *ICmp, Value *X, bool IsAnd) {
+  // icmp eq/ne (X & mask), 0
+  ICmpInst::Predicate Pred;
+  const APInt *Mask;
+  if (match(ICmp,
+            m_ICmp(Pred, m_And(m_Specific(X), m_APInt(Mask)), m_Zero())) &&
+      ICmp->isEquality()) {
+    if (Mask->popcount() == Mask->getBitWidth() - 1) {
+      auto Zero = APInt::getZero(Mask->getBitWidth());
+      auto Val = ~*Mask;
+      if (IsAnd) {
+        if (Pred == ICmpInst::ICMP_EQ) {
+          Set.push_back(ConstantRange(Zero, Val + 1));
+          Set.push_back(ConstantRange(Val, APInt(Mask->getBitWidth(), 1)));
+        } else {
+          Set.push_back(ConstantRange(Zero).inverse());
+          Set.push_back(ConstantRange(Val).inverse());
+        }
+      } else {
+        if (Pred == ICmpInst::ICMP_EQ) {
+          Set.push_back(Val);
+          Set.push_back(Zero);
+        } else {
+          Set.push_back(ConstantRange(APInt(Mask->getBitWidth(), 1), Val));
+          Set.push_back(ConstantRange(Val + 1, Zero));
+        }
+      }
+      return true;
+    }
+  }
+
+  return false;
+}
+
+/// Fold (icmp Pred1 V1, C1) & (icmp Pred2 V2, C2)
+/// or   (icmp Pred1 V1, C1) | (icmp Pred2 V2, C2)
+/// into a single comparison using range-based reasoning.
+/// It handles patterns which cannot be recognized by
+/// foldAndOrOfICmpsUsingRanges. Try to decompose one of icmps into two or more
+/// icmps. Example: icmp eq (X & ~signmask), 0 --> (icmp eq X, 0) | (icmp eq X,
+/// signmask)
+static Value *
+foldAndOrOfICmpsUsingDecomposedRanges(ICmpInst *ICmp1, ICmpInst *ICmp2,
+                                      bool IsAnd,
+                                      InstCombiner::BuilderTy &Builder) {
+  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;
+
+  SmallVector<ConstantRange, 3> Ranges;
+  Value *X = nullptr;
+  if (decomposeICmpIntoRangeSet(Ranges, ICmp2, V1, IsAnd)) {
+    X = V1;
+    Ranges.push_back(ConstantRange::makeExactICmpRegion(Pred1, *C1));
+  } else if (decomposeICmpIntoRangeSet(Ranges, ICmp1, V2, IsAnd)) {
+    X = V2;
+    Ranges.push_back(ConstantRange::makeExactICmpRegion(Pred2, *C2));
+  } else
+    return nullptr;
+
+  while (true) {
+    bool Merged = false;
+    for (unsigned I = 0; I < Ranges.size(); ++I) {
+      auto &CR = Ranges[I];
+      for (unsigned J = I + 1; J < Ranges.size(); ++J) {
+        if (auto NewCR = IsAnd ? CR.exactIntersectWith(Ranges[J])
+                               : CR.exactUnionWith(Ranges[J])) {
+          CR = *NewCR;
+          Ranges.erase(Ranges.begin() + J);
+          Merged = true;
+        }
+      }
+    }
+
+    if (Ranges.size() == 1)
+      break;
+
+    if (!Merged)
+      return nullptr;
+  }
+
+  ICmpInst::Predicate NewPred;
+  APInt NewRHS, NewOffset;
+  Ranges[0].getEquivalentICmp(NewPred, NewRHS, NewOffset);
+  // Similar to foldAndOrOfICmpsUsingRanges, we don't check hasOneUse here.
+  if (NewOffset != 0)
+    X = Builder.CreateAdd(X, ConstantInt::get(X->getType(), NewOffset));
+  return Builder.CreateICmp(NewPred, X, ConstantInt::get(X->getType(), NewRHS));
+}
+
 /// Ignore all operations which only change the sign of a value, returning the
 /// underlying magnitude value.
 static Value *stripSignOnlyFPOps(Value *Val) {
@@ -3280,7 +3375,9 @@ Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
     }
   }
 
-  return foldAndOrOfICmpsUsingRanges(LHS, RHS, IsAnd);
+  if (auto *V = foldAndOrOfICmpsUsingRanges(LHS, RHS, IsAnd))
+    return V;
+  return foldAndOrOfICmpsUsingDecomposedRanges(LHS, RHS, IsAnd, Builder);
 }
 
 // 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 7af06e03fd4b2a9..1191b553c361673 100644
--- a/llvm/test/Transforms/InstCombine/icmp-range.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-range.ll
@@ -1521,6 +1521,112 @@ define i1 @isFloat(i64 %0) {
   ret i1 %5
 }
 
+; tests from PR69123
+define i1 @or_slt_eq_masked(i32 %a) {
+; CHECK-LABEL: @or_slt_eq_masked(
+; CHECK-NEXT:    [[OR:%.*]] = icmp slt i32 [[A:%.*]], 1
+; CHECK-NEXT:    ret i1 [[OR]]
+;
+  %cmp = icmp slt i32 %a, 0
+  %and1 = and i32 %a, 2147483647
+  %tobool2.not = icmp eq i32 %and1, 0
+  %or = or i1 %cmp, %tobool2.not
+  ret i1 %or
+}
+
+define i1 @and_sgt_ne_masked(i32 %d) {
+; CHECK-LABEL: @and_sgt_ne_masked(
+; CHECK-NEXT:    [[AND:%.*]] = icmp sgt i32 [[D:%.*]], 0
+; CHECK-NEXT:    ret i1 [[AND]]
+;
+  %tobool.not = icmp sgt i32 %d, -1
+  %and1 = and i32 %d, 2147483647
+  %tobool2 = icmp ne i32 %and1, 0
+  %and = and i1 %tobool.not, %tobool2
+  ret i1 %and
+}
+
+define i1 @or_slt_ne_masked(i32 %a) {
+; CHECK-LABEL: @or_slt_ne_masked(
+; CHECK-NEXT:    [[OR:%.*]] = icmp ne i32 [[A:%.*]], 0
+; CHECK-NEXT:    ret i1 [[OR]]
+;
+  %cmp = icmp slt i32 %a, 0
+  %and1 = and i32 %a, 2147483647
+  %tobool2.not = icmp ne i32 %and1, 0
+  %or = or i1 %cmp, %tobool2.not
+  ret i1 %or
+}
+
+define i1 @and_sgt_eq_masked(i32 %d) {
+; CHECK-LABEL: @and_sgt_eq_masked(
+; CHECK-NEXT:    [[AND:%.*]] = icmp eq i32 [[D:%.*]], 0
+; CHECK-NEXT:    ret i1 [[AND]]
+;
+  %tobool.not = icmp sgt i32 %d, -1
+  %and1 = and i32 %d, 2147483647
+  %tobool2 = icmp eq i32 %and1, 0
+  %and = and i1 %tobool.not, %tobool2
+  ret i1 %and
+}
+
+define i1 @or_slt_eq_masked2(i32 %a) {
+; CHECK-LABEL: @or_slt_eq_masked2(
+; CHECK-NEXT:    [[OR:%.*]] = icmp slt i32 [[A:%.*]], 9
+; CHECK-NEXT:    ret i1 [[OR]]
+;
+  %cmp = icmp slt i32 %a, 8
+  %and1 = and i32 %a, -9
+  %tobool2.not = icmp eq i32 %and1, 0
+  %or = or i1 %cmp, %tobool2.not
+  ret i1 %or
+}
+
+define i1 @or_slt_eq_masked_nofold_invalid_mask(i32 %a) {
+; CHECK-LABEL: @or_slt_eq_masked_nofold_invalid_mask(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[A]], 2147483646
+; CHECK-NEXT:    [[TOBOOL2_NOT:%.*]] = icmp eq i32 [[AND1]], 0
+; CHECK-NEXT:    [[OR:%.*]] = or i1 [[CMP]], [[TOBOOL2_NOT]]
+; CHECK-NEXT:    ret i1 [[OR]]
+;
+  %cmp = icmp slt i32 %a, 0
+  %and1 = and i32 %a, 2147483646
+  %tobool2.not = icmp eq i32 %and1, 0
+  %or = or i1 %cmp, %tobool2.not
+  ret i1 %or
+}
+
+define i1 @or_slt_eq_masked_nofold_unmergeable_ranges(i32 %a) {
+; CHECK-LABEL: @or_slt_eq_masked_nofold_unmergeable_ranges(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[A:%.*]], -1
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[A]], 2147483647
+; CHECK-NEXT:    [[TOBOOL2_NOT:%.*]] = icmp eq i32 [[AND1]], 0
+; CHECK-NEXT:    [[OR:%.*]] = or i1 [[CMP]], [[TOBOOL2_NOT]]
+; CHECK-NEXT:    ret i1 [[OR]]
+;
+  %cmp = icmp slt i32 %a, -1
+  %and1 = and i32 %a, 2147483647
+  %tobool2.not = icmp eq i32 %and1, 0
+  %or = or i1 %cmp, %tobool2.not
+  ret i1 %or
+}
+
+define i1 @and_sgt_ne_masked_nofold_unmergeable_ranges(i32 %d) {
+; CHECK-LABEL: @and_sgt_ne_masked_nofold_unmergeable_ranges(
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp sgt i32 [[D:%.*]], -2
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[D]], 2147483647
+; CHECK-NEXT:    [[TOBOOL2:%.*]] = icmp ne i32 [[AND1]], 0
+; CHECK-NEXT:    [[AND:%.*]] = and i1 [[TOBOOL_NOT]], [[TOBOOL2]]
+; CHECK-NEXT:    ret i1 [[AND]]
+;
+  %tobool.not = icmp sgt i32 %d, -2
+  %and1 = and i32 %d, 2147483647
+  %tobool2 = icmp ne i32 %and1, 0
+  %and = and i1 %tobool.not, %tobool2
+  ret i1 %and
+}
+
 !0 = !{i32 1, i32 6}
 !1 = !{i32 0, i32 6}
 !2 = !{i8 0, i8 1}

[dtcxzyw]

Ping.