[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 da5f625ebb1f5dd0630c1197e79d02566f795a5c 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 f9d59a7ed64cd6f17cdc22d2706cce6b0612b35f 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] [InstCombine] Canonicalize (a + 1 == 0) ? -1 : a + 1 ->
 uadd.sat(a, 1)

This is because this actually is just:

// ((X + Y) u< X) ? -1 : (X + Y) --> uadd.sat(X, Y)
// ((X + Y) u< Y) ? -1 : (X + Y) --> uadd.sat(X, Y)

However, ult 1 is canonicalized to eq 0.

Alive2: https://alive2.llvm.org/ce/z/Y4eHav
---
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp | 11 +++++++++--
 .../test/Transforms/InstCombine/saturating-add-sub.ll |  4 +---
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 73ba0f78e8053..d733a37efd68e 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1013,12 +1013,19 @@ static Value *canonicalizeSaturatedAdd(ICmpInst *Cmp, Value *TVal, Value *FVal,
 
   // uge -1 is canonicalized to eq -1 and requires special handling
   // (a == -1) ? -1 : a + 1 -> uadd.sat(a, 1)
+  // ult 1 is canonicalized to eq 0
+  // (a + 1 == 0) ? -1 : a + 1 -> uadd.sat(a, 1)
   if (Pred == ICmpInst::ICMP_EQ) {
     if (match(FVal, m_Add(m_Specific(Cmp0), m_One())) &&
-        match(Cmp1, m_AllOnes())) {
+        match(Cmp1, m_AllOnes()))
       return Builder.CreateBinaryIntrinsic(
           Intrinsic::uadd_sat, Cmp0, ConstantInt::get(Cmp0->getType(), 1));
-    }
+
+    if (match(Cmp1, m_Zero()) && match(Cmp0, m_Add(m_Value(X), m_One())) &&
+        match(FVal, m_Add(m_Specific(X), m_One())))
+      return Builder.CreateBinaryIntrinsic(Intrinsic::uadd_sat, X,
+                                           ConstantInt::get(X->getType(), 1));
+
     return nullptr;
   }
 
diff --git a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
index b3bad2f88288a..392551defb7ef 100644
--- a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
+++ b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
@@ -2352,9 +2352,7 @@ define i8 @fold_add_umax_to_usub_multiuse(i8 %a) {
 
 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:    [[COND:%.*]] = call i32 @llvm.uadd.sat.i32(i32 [[ADD:%.*]], i32 1)
 ; CHECK-NEXT:    ret i32 [[COND]]
 ;
   %add = add i32 %num, 1