[llvm] [InstCombine] Fix saturated add canonicalization (PR #97973)

[via llvm-commits]

https://github.com/AtariDreams updated https://github.com/llvm/llvm-project/pull/97973

>From f42eecaecfa7a97361ec8913e8713d736a32e512 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sun, 7 Jul 2024 17:51:02 -0400
Subject: [PATCH 1/2] Add flipped version of (X u< ~C) ? (X + C) : -1 -->
 uadd.sat(X, C)

For some reason, LLVM is not considering X u > C, a, b to be the same as X <= C, b, a.

Move that to after the canonicalization.
---
 .../InstCombine/InstCombineSelect.cpp         | 24 +++++++++++++------
 .../InstCombine/saturating-add-sub.ll         | 12 ++++++++++
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 394dfca262e130..6958f890a5fa5f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -978,13 +978,6 @@ static Value *canonicalizeSaturatedAdd(ICmpInst *Cmp, Value *TVal, Value *FVal,
   ICmpInst::Predicate Pred = Cmp->getPredicate();
   Value *X;
   const APInt *C, *CmpC;
-  if (Pred == ICmpInst::ICMP_ULT &&
-      match(TVal, m_Add(m_Value(X), m_APInt(C))) && X == Cmp0 &&
-      match(FVal, m_AllOnes()) && match(Cmp1, m_APInt(CmpC)) && *CmpC == ~*C) {
-    // (X u< ~C) ? (X + C) : -1 --> uadd.sat(X, C)
-    return Builder.CreateBinaryIntrinsic(
-        Intrinsic::uadd_sat, X, ConstantInt::get(X->getType(), *C));
-  }
 
   // Match unsigned saturated add of 2 variables with an unnecessary 'not'.
   // There are 8 commuted variants.
@@ -996,6 +989,23 @@ static Value *canonicalizeSaturatedAdd(ICmpInst *Cmp, Value *TVal, Value *FVal,
   if (!match(TVal, m_AllOnes()))
     return nullptr;
 
+  if (Pred == ICmpInst::ICMP_UGE &&
+      match(FVal, m_Add(m_Value(X), m_APInt(C))) && X == Cmp0 &&
+      match(TVal, m_AllOnes()) && match(Cmp1, m_APInt(CmpC)) && *CmpC == ~*C) {
+    // (X u >= ~C) ? -1 : (X + C) --> uadd.sat(X, C)
+    return Builder.CreateBinaryIntrinsic(Intrinsic::uadd_sat, X,
+                                         ConstantInt::get(X->getType(), *C));
+  }
+
+  if (Pred == ICmpInst::ICMP_UGT &&
+      match(FVal, m_Add(m_Value(X), m_APInt(C))) && X == Cmp0 &&
+      match(TVal, m_AllOnes()) && match(Cmp1, m_APInt(CmpC)) &&
+      *CmpC == (~*C - 1)) {
+    // (X u > ~C - 1) ? -1 : (X + C) --> uadd.sat(X, C)
+    return Builder.CreateBinaryIntrinsic(Intrinsic::uadd_sat, X,
+                                         ConstantInt::get(X->getType(), *C));
+  }
+
   // Canonicalize predicate to less-than or less-or-equal-than.
   if (Pred == ICmpInst::ICMP_UGT || Pred == ICmpInst::ICMP_UGE) {
     std::swap(Cmp0, Cmp1);
diff --git a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
index bf1568f1cd8c09..1db6277c10d023 100644
--- a/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
+++ b/llvm/test/Transforms/InstCombine/saturating-add-sub.ll
@@ -1395,6 +1395,18 @@ define i32 @uadd_sat(i32 %x, i32 %y) {
   %r = select i1 %c, i32 -1, i32 %a
   ret i32 %r
 }
+
+define i32 @uadd_sat_flipped(i32 %x) {
+; CHECK-LABEL: @uadd_sat_flipped(
+; CHECK-NEXT:    [[COND:%.*]] = call i32 @llvm.uadd.sat.i32(i32 [[X:%.*]], i32 9)
+; CHECK-NEXT:    ret i32 [[COND]]
+;
+  %cmp = icmp uge i32 %x, -10
+  %add = add i32 %x, 9
+  %cond = select i1 %cmp, i32 -1, i32 %add
+  ret i32 %cond
+}
+
 define i32 @uadd_sat_nonstrict(i32 %x, i32 %y) {
 ; CHECK-LABEL: @uadd_sat_nonstrict(
 ; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.uadd.sat.i32(i32 [[X:%.*]], i32 [[Y:%.*]])

>From 808a49adfe1ec2fa37c959401cef0dfb74e20cee Mon Sep 17 00:00:00 2001
From: AtariDreams <gfunni234 at gmail.com>
Date: Mon, 8 Jul 2024 17:03:05 -0400
Subject: [PATCH 2/2] Update
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Co-authored-by: Nikita Popov <github at npopov.com>
---
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 6958f890a5fa5f..61cf8a38a53429 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -992,7 +992,7 @@ static Value *canonicalizeSaturatedAdd(ICmpInst *Cmp, Value *TVal, Value *FVal,
   if (Pred == ICmpInst::ICMP_UGE &&
       match(FVal, m_Add(m_Value(X), m_APInt(C))) && X == Cmp0 &&
       match(TVal, m_AllOnes()) && match(Cmp1, m_APInt(CmpC)) && *CmpC == ~*C) {
-    // (X u >= ~C) ? -1 : (X + C) --> uadd.sat(X, C)
+    // (X u>= ~C) ? -1 : (X + C) --> uadd.sat(X, C)
     return Builder.CreateBinaryIntrinsic(Intrinsic::uadd_sat, X,
                                          ConstantInt::get(X->getType(), *C));
   }