[llvm] [InstCombine] Canonicalize (a + 1 == 0) ? -1 : a + 1 -> uadd.sat(a, 1) (PR #144566)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/144566

>From 3382338dbe6eb691fff2828462162657256df956 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Tue, 17 Jun 2025 12:40:21 -0400
Subject: [PATCH 1/2] [InstCombine] Pre-commit test

---
 .../Transforms/InstCombine/saturating-add-sub.ll    | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
index cfd679c0cc592..b3bad2f88288a 100644
--- a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
+++ b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
@@ -2350,4 +2350,17 @@ define i8 @fold_add_umax_to_usub_multiuse(i8 %a) {
   ret i8 %sel
 }
 
+define i32 @add_check_zero(i32 %num) {
+; CHECK-LABEL: @add_check_zero(
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[NUM:%.*]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[ADD]], 0
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i32 -1, i32 [[ADD]]
+; CHECK-NEXT:    ret i32 [[COND]]
+;
+  %add = add i32 %num, 1
+  %cmp = icmp eq i32 %add, 0
+  %cond = select i1 %cmp, i32 -1, i32 %add
+  ret i32 %cond
+}
+
 declare void @usei8(i8)

>From 932b5aa731f4fc3e2b647e804898e858c096ce39 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Tue, 17 Jun 2025 13:01:36 -0400
Subject: [PATCH 2/2] Allow folding icmp eq (add X, C2), C when there is more
 than one-use when we can compute the range

If there are multiple uses of an add, we can fold a comparison anyway if we can compute a constant range for it, which should happen in cases such as saturated add.
---
 llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp | 7 +++++--
 llvm/test/Transforms/InstCombine/uaddo.ll               | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 084e7fbaa268a..e6542223d6956 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -3160,7 +3160,7 @@ Instruction *InstCombinerImpl::foldICmpAddConstant(ICmpInst &Cmp,
       return replaceInstUsesWith(Cmp, Cond);
   }
   const APInt *C2;
-  if (Cmp.isEquality() || !match(Y, m_APInt(C2)))
+  if (!match(Y, m_APInt(C2)))
     return nullptr;
 
   // Fold icmp pred (add X, C2), C.
@@ -3184,7 +3184,7 @@ Instruction *InstCombinerImpl::foldICmpAddConstant(ICmpInst &Cmp,
       return new ICmpInst(Pred, X, ConstantInt::get(Ty, NewC));
   }
 
-  if (ICmpInst::isUnsigned(Pred) && Add->hasNoSignedWrap() &&
+  if (!Cmp.isEquality() && ICmpInst::isUnsigned(Pred) && Add->hasNoSignedWrap() &&
       C.isNonNegative() && (C - *C2).isNonNegative() &&
       computeConstantRange(X, /*ForSigned=*/true).add(*C2).isAllNonNegative())
     return new ICmpInst(ICmpInst::getSignedPredicate(Pred), X,
@@ -3205,6 +3205,9 @@ Instruction *InstCombinerImpl::foldICmpAddConstant(ICmpInst &Cmp,
       return new ICmpInst(ICmpInst::ICMP_UGE, X, ConstantInt::get(Ty, Lower));
   }
 
+  if (Cmp.isEquality())
+    return nullptr;
+
   // This set of folds is intentionally placed after folds that use no-wrapping
   // flags because those folds are likely better for later analysis/codegen.
   const APInt SMax = APInt::getSignedMaxValue(Ty->getScalarSizeInBits());
diff --git a/llvm/test/Transforms/InstCombine/uaddo.ll b/llvm/test/Transforms/InstCombine/uaddo.ll
index ae7a07ec8000c..89a9569bdb7c3 100644
--- a/llvm/test/Transforms/InstCombine/uaddo.ll
+++ b/llvm/test/Transforms/InstCombine/uaddo.ll
@@ -158,7 +158,7 @@ define i1 @uaddo_1(i8 %x, ptr %p) {
 ; CHECK-LABEL: @uaddo_1(
 ; CHECK-NEXT:    [[A:%.*]] = add i8 [[X:%.*]], 1
 ; CHECK-NEXT:    store i8 [[A]], ptr [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[A]], 0
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[X]], -1
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %a = add i8 %x, 1