[llvm] [InstCombine] Fold (zext (X +nuw C)) + -C --> zext(X) when zext has additional use. (PR #98533)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/98533

>From 971e930cd093e1d08ed9002ef9be617f64d3131e Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 11 Jul 2024 12:56:01 -0700
Subject: [PATCH 1/2] [InstCombine] Add additional test case for (add (zext
 (add nuw X, C2), C1). NFC

Add test where the zext has an additional use, but the entire
expression can be replaced with (zext X). Folding even though there
is an additional use would not increase the number of instructions.
---
 llvm/test/Transforms/InstCombine/add.ll | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/add.ll b/llvm/test/Transforms/InstCombine/add.ll
index 8b3ecaf25cd9d..2f914a2e58867 100644
--- a/llvm/test/Transforms/InstCombine/add.ll
+++ b/llvm/test/Transforms/InstCombine/add.ll
@@ -951,6 +951,21 @@ define i64 @test41(i32 %a) {
   ret i64 %sub
 }
 
+define i64 @test41_multiuse_constants_cancel(i32 %a) {
+; CHECK-LABEL: @test41_multiuse_constants_cancel(
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i32 [[A:%.*]], 1
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[REASS_ADD:%.*]] = shl nuw nsw i64 [[ZEXT]], 1
+; CHECK-NEXT:    [[EXTRAUSE:%.*]] = add nsw i64 [[REASS_ADD]], -1
+; CHECK-NEXT:    ret i64 [[EXTRAUSE]]
+;
+  %add = add nuw i32 %a, 1
+  %zext = zext i32 %add to i64
+  %sub = add i64 %zext, -1
+  %extrause = add i64 %zext, %sub
+  ret i64 %extrause
+}
+
 ; (add (zext (add nuw X, C2)), C) --> (zext (add nuw X, C2 + C))
 
 define <2 x i64> @test41vec(<2 x i32> %a) {

>From 42fdcf8af82576a8f56e2d255e2a193e68a9e324 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 11 Jul 2024 12:59:29 -0700
Subject: [PATCH 2/2] [InstCombine] Fold (zext (X +nuw C)) + -C --> zext(X)
 when zext has additional use.

We have a general fold for (zext (X +nuw C2)) + C1 --> zext (X + (C2 + trunc(C1)))
but this fold is disabled if the zext has an additional use.

If the two constants cancel, we can fold the whole expression to
zext(X) without increasing the number of instructions.

Though we will have 2 zexts which might be not be cheap. So
another option could be to move the original narrow add after the zext
for the other uses. That would allow sharing of the zext(x).
---
 .../Transforms/InstCombine/InstCombineAddSub.cpp    | 13 +++++++++----
 llvm/test/Transforms/InstCombine/add.ll             |  4 ++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index 0a73c58c07409..0a55f4762fdf0 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -819,11 +819,16 @@ static Instruction *foldNoWrapAdd(BinaryOperator &Add,
   Value *X;
   const APInt *C1, *C2;
   if (match(Op1, m_APInt(C1)) &&
-      match(Op0, m_OneUse(m_ZExt(m_NUWAddLike(m_Value(X), m_APInt(C2))))) &&
+      match(Op0, m_ZExt(m_NUWAddLike(m_Value(X), m_APInt(C2)))) &&
       C1->isNegative() && C1->sge(-C2->sext(C1->getBitWidth()))) {
-    Constant *NewC =
-        ConstantInt::get(X->getType(), *C2 + C1->trunc(C2->getBitWidth()));
-    return new ZExtInst(Builder.CreateNUWAdd(X, NewC), Ty);
+    APInt NewC = *C2 + C1->trunc(C2->getBitWidth());
+    // If the smaller add will fold to zero, we don't need to check one use.
+    if (NewC.isZero())
+      return new ZExtInst(X, Ty);
+    // Otherwise only do this if the existing zero extend will be removed.
+    if (Op0->hasOneUse())
+      return new ZExtInst(
+          Builder.CreateNUWAdd(X, ConstantInt::get(X->getType(), NewC)), Ty);
   }
 
   // More general combining of constants in the wide type.
diff --git a/llvm/test/Transforms/InstCombine/add.ll b/llvm/test/Transforms/InstCombine/add.ll
index 2f914a2e58867..adb61cd4b9692 100644
--- a/llvm/test/Transforms/InstCombine/add.ll
+++ b/llvm/test/Transforms/InstCombine/add.ll
@@ -955,8 +955,8 @@ define i64 @test41_multiuse_constants_cancel(i32 %a) {
 ; CHECK-LABEL: @test41_multiuse_constants_cancel(
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw i32 [[A:%.*]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 [[ADD]] to i64
-; CHECK-NEXT:    [[REASS_ADD:%.*]] = shl nuw nsw i64 [[ZEXT]], 1
-; CHECK-NEXT:    [[EXTRAUSE:%.*]] = add nsw i64 [[REASS_ADD]], -1
+; CHECK-NEXT:    [[SUB:%.*]] = zext i32 [[A]] to i64
+; CHECK-NEXT:    [[EXTRAUSE:%.*]] = add nuw nsw i64 [[ZEXT]], [[SUB]]
 ; CHECK-NEXT:    ret i64 [[EXTRAUSE]]
 ;
   %add = add nuw i32 %a, 1