[llvm] [InstCombine] Decompose an icmp into multiple ranges (PR #69855)

[via llvm-commits]

================
@@ -1334,6 +1334,122 @@ 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) {
+  // Handle "icmp eq/ne (X & mask), 0", where mask is an inverted power of 2.
+  ICmpInst::Predicate Pred;
+  const APInt *Mask;
+  if (match(ICmp,
+            m_ICmp(Pred, m_And(m_Specific(X), m_APInt(Mask)), m_Zero())) &&
+      ICmp->isEquality() && (Mask->popcount() == Mask->getBitWidth() - 1)) {
+    auto Zero = APInt::getZero(Mask->getBitWidth());
+    auto Val = ~*Mask;
+    if (IsAnd) {
+      if (Pred == ICmpInst::ICMP_EQ) {
+        // icmp eq (X & mask), 0
+        // X in {0, Val} --> X in [0, Val] and X in [Val, 0] (upper-wrapped)
+        Set.push_back(ConstantRange(Zero, Val + 1));
+        Set.push_back(ConstantRange(Val, APInt(Mask->getBitWidth(), 1)));
+      } else {
+        // icmp ne (X & mask), 0
+        // X not in {0, Val} --> X not in {0} and X not in {Val}
+        Set.push_back(ConstantRange(Zero).inverse());
+        Set.push_back(ConstantRange(Val).inverse());
+      }
+    } else {
+      if (Pred == ICmpInst::ICMP_EQ) {
+        // icmp eq (X & mask), 0
+        // X in {0, Val} --> X in {0} or X in {Val}
+        Set.push_back(Val);
+        Set.push_back(Zero);
+      } else {
+        // icmp ne (X & mask), 0
+        // X not in {0, Val} --> X in [1, Val) or X not in [Val + 1, 0)
+        // (upper-wrapped)
+        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;
+
+  // Try to merge ranges into single range.
+  // Since we may fail to merge ranges due to the order of merging, we cannot do
+  // merge in order (counterexample: [0, 1), [2, 3), [1, 2)) or in sorted order
+  // (due to wrapped ranges). Instead, we try to merge ranges in O(N^2) until we
+  // succeed in merging into single range or we cannot merge anymore.
+
+  // Ranges which cannot be merged with each other. We maintain this list to
+  // avoid redundant checks.
+  SmallVector<ConstantRange, 3> UnmergeableRanges;
+  while (true) {
+    bool Merged = false;
+
+    if (!Ranges.empty()) {
+      auto &CR = Ranges.back();
+      for (unsigned I = 0; I < UnmergeableRanges.size(); ++I) {
+        if (auto NewCR = IsAnd ? CR.exactIntersectWith(UnmergeableRanges[I])
+                               : CR.exactUnionWith(UnmergeableRanges[I])) {
+          CR = *NewCR;
+          UnmergeableRanges.erase(UnmergeableRanges.begin() + I);
+          Merged = true;
+        }
+      }
+      if (!Merged) {
+        UnmergeableRanges.push_back(CR);
+        Ranges.pop_back();
+      }
+    }
+
+    if (Ranges.empty()) {
+      if (UnmergeableRanges.size() == 1)
+        break;
+      return nullptr;
----------------
goldsteinn wrote:

Stylistically would just break on `Range.empty()` and move to `size() != 1 { return nullptr; }` to outside the loop.

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