[llvm] [InstCombine] Add optimization to combine adds through zext nneg, and add testcase (PR #157723)

[via llvm-commits]

https://github.com/Aethezz updated https://github.com/llvm/llvm-project/pull/157723

>From 0e8a5fd5b05517159b0df102913086d6e4bb2390 Mon Sep 17 00:00:00 2001
From: Aethezz <ellisonlao999 at gmail.com>
Date: Tue, 9 Sep 2025 13:20:22 -0400
Subject: [PATCH 1/2] Combine adds through zext nneg, and add testcase

---
 .../InstCombine/InstCombineAddSub.cpp           | 17 +++++++++++++++++
 llvm/test/Transforms/InstCombine/zext.ll        | 11 +++++++++++
 2 files changed, 28 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index d934638c15e75..f4a2ab5377a34 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -1910,6 +1910,23 @@ Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
   if (Instruction *Res = foldBinOpOfSelectAndCastOfSelectCondition(I))
     return Res;
 
+  // Combine adds through zext nneg:
+  // (zext nneg (X + C1)) + C2 --> zext X if C1 + C2 == 0
+  {
+    Value *X;
+    const APInt *C1, *C2;
+    if (match(&I, m_c_Add(m_NNegZExt(m_Add(m_Value(X), m_APInt(C1))),
+                          m_APInt(C2)))) {
+      // Check if the constants cancel out (C1 + C2 == 0)
+      APInt Sum = C1->sext(C2->getBitWidth()) + *C2;
+      if (Sum.isZero()) {
+        // The add inside the zext and the outer add cancel out
+        Value *NewZExt = Builder.CreateZExt(X, I.getType());
+        return replaceInstUsesWith(I, NewZExt);
+      }
+    }
+  }
+
   // Re-enqueue users of the induction variable of add recurrence if we infer
   // new nuw/nsw flags.
   if (Changed) {
diff --git a/llvm/test/Transforms/InstCombine/zext.ll b/llvm/test/Transforms/InstCombine/zext.ll
index e4d18e9395219..15aa2aaec430a 100644
--- a/llvm/test/Transforms/InstCombine/zext.ll
+++ b/llvm/test/Transforms/InstCombine/zext.ll
@@ -976,3 +976,14 @@ entry:
   %res = zext nneg i2 %x to i32
   ret i32 %res
 }
+
+define i32 @zext_nneg_add_cancel(i8 %arg) {
+; CHECK-LABEL: @zext_nneg_add_cancel(
+; CHECK-NEXT:    [[ADD2:%.*]] = zext i8 [[ARG:%.*]] to i32
+; CHECK-NEXT:    ret i32 [[ADD2]]
+;
+  %add = add i8 %arg, -2
+  %zext = zext nneg i8 %add to i32
+  %add2 = add i32 %zext, 2
+  ret i32 %add2
+}
\ No newline at end of file

>From 33305609276cb3aaef7f66766372125c7e7fd589 Mon Sep 17 00:00:00 2001
From: Aethezz <ellisonlao999 at gmail.com>
Date: Tue, 16 Sep 2025 12:48:02 -0400
Subject: [PATCH 2/2] added preconditions and adjusted testcases

---
 .../InstCombine/InstCombineAddSub.cpp         | 28 +++++++++++++------
 llvm/test/Transforms/InstCombine/zext.ll      | 27 +++++++++++++++++-
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index f4a2ab5377a34..ffecb558e332a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -1910,20 +1910,32 @@ Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
   if (Instruction *Res = foldBinOpOfSelectAndCastOfSelectCondition(I))
     return Res;
 
-  // Combine adds through zext nneg:
-  // (zext nneg (X + C1)) + C2 --> zext X if C1 + C2 == 0
   {
     Value *X;
     const APInt *C1, *C2;
     if (match(&I, m_c_Add(m_NNegZExt(m_Add(m_Value(X), m_APInt(C1))),
                           m_APInt(C2)))) {
-      // Check if the constants cancel out (C1 + C2 == 0)
+      // Check if inner constant C1 is negative C1 < 0
+      if (!C1->isNegative())
+        return nullptr;
+
+      if (C2->isNegative())
+        return nullptr; // We only handle C2 >= 0
+
       APInt Sum = C1->sext(C2->getBitWidth()) + *C2;
-      if (Sum.isZero()) {
-        // The add inside the zext and the outer add cancel out
-        Value *NewZExt = Builder.CreateZExt(X, I.getType());
-        return replaceInstUsesWith(I, NewZExt);
-      }
+      APInt newSum = Sum.trunc(C1->getBitWidth());
+
+      if (newSum.sext(C2->getBitWidth()) != Sum)
+        return nullptr;
+
+      // X if sum is zero, else X + newSum
+      Value *Inner =
+          Sum.isZero()
+              ? X
+              : Builder.CreateAdd(X, ConstantInt::get(X->getType(), newSum));
+
+      Value *NewZExt = Builder.CreateZExt(Inner, I.getType());
+      return replaceInstUsesWith(I, NewZExt);
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/zext.ll b/llvm/test/Transforms/InstCombine/zext.ll
index 15aa2aaec430a..c2622aab437af 100644
--- a/llvm/test/Transforms/InstCombine/zext.ll
+++ b/llvm/test/Transforms/InstCombine/zext.ll
@@ -986,4 +986,29 @@ define i32 @zext_nneg_add_cancel(i8 %arg) {
   %zext = zext nneg i8 %add to i32
   %add2 = add i32 %zext, 2
   ret i32 %add2
-}
\ No newline at end of file
+}
+
+define i32 @zext_nneg_no_cancel(i8 %arg) {
+; CHECK-LABEL: @zext_nneg_no_cancel(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[ARG:%.*]], -1
+; CHECK-NEXT:    [[ADD2:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[ADD2]]
+;
+  %add = add i8 %arg, -2
+  %zext = zext nneg i8 %add to i32
+  %add2 = add i32 %zext, 1
+  ret i32 %add2
+}
+
+define i32 @zext_nneg_overflow(i8 %arg) {
+; CHECK-LABEL: @zext_nneg_overflow(
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[ARG:%.*]], -2
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext nneg i8 [[ADD]] to i32
+; CHECK-NEXT:    [[ADD2:%.*]] = add nuw nsw i32 [[ZEXT]], 299
+; CHECK-NEXT:    ret i32 [[ADD2]]
+;
+  %add = add i8 %arg, -2
+  %zext = zext nneg i8 %add to i32
+  %add2 = add i32 %zext, 299
+  ret i32 %add2
+}