[llvm] [InstCombine] Decompose icmp eq/ne with add/sub trees (PR #92881)

[Yingwei Zheng via llvm-commits]

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

Motivating case: https://alive2.llvm.org/ce/z/8S6dRK

>From dc7fa07cd4abac12a65163e6aeb1a01dfa02107e Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Tue, 21 May 2024 16:42:49 +0800
Subject: [PATCH] [InstCombine] Decompose icmp eq/ne with add/sub trees

---
 .../InstCombine/InstCombineCompares.cpp       | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 542a1c82b127a..01bdb52499af5 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -5447,6 +5447,53 @@ static Instruction *foldICmpPow2Test(ICmpInst &I,
   return nullptr;
 }
 
+static void decomposeAddSub(Value *V, DenseMap<Value *, int32_t> &Row,
+                            Constant *&C, uint32_t Depth, bool Add) {
+  if (auto *I = dyn_cast<Instruction>(V)) {
+    if (Depth && !Row.contains(V) &&
+        (I->getOpcode() == Instruction::Add ||
+         I->getOpcode() == Instruction::Sub)) {
+      decomposeAddSub(I->getOperand(0), Row, C, Depth - 1, Add);
+      decomposeAddSub(I->getOperand(1), Row, C, Depth - 1,
+                      (I->getOpcode() == Instruction::Sub) ^ Add);
+      return;
+    }
+  }
+
+  Constant *CI;
+  if (match(V, m_ImmConstant(CI))) {
+    if (auto *Res =
+            (Add ? ConstantExpr::getAdd(C, CI) : ConstantExpr::getSub(C, CI))) {
+      C = Res;
+      return;
+    }
+  }
+
+  Row[V] += Add ? 1 : -1;
+}
+
+static Value *foldICmpEqualityWithAddSubTree(ICmpInst::Predicate Pred,
+                                             Value *Op0, Value *Op1,
+                                             const SimplifyQuery &SQ) {
+  if (Op0->getType()->isPtrOrPtrVectorTy())
+    return nullptr;
+
+  if (match(Op1, m_Constant()))
+    return nullptr;
+
+  DenseMap<Value *, int32_t> Row;
+  Constant *C = Constant::getNullValue(Op0->getType());
+  constexpr uint32_t Depth = 4;
+  decomposeAddSub(Op0, Row, C, Depth, /*Add=*/true);
+  decomposeAddSub(Op1, Row, C, Depth, /*Add=*/false);
+
+  for (auto &[V, K] : Row)
+    if (K != 0)
+      return nullptr;
+
+  return ConstantExpr::getICmp(Pred, C, Constant::getNullValue(Op0->getType()));
+}
+
 Instruction *InstCombinerImpl::foldICmpEquality(ICmpInst &I) {
   if (!I.isEquality())
     return nullptr;
@@ -5695,6 +5742,10 @@ Instruction *InstCombinerImpl::foldICmpEquality(ICmpInst &I) {
                                   : ConstantInt::getNullValue(A->getType()));
   }
 
+  if (auto *V = foldICmpEqualityWithAddSubTree(Pred, Op0, Op1,
+                                               SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, V);
+
   return nullptr;
 }