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

[topperc]

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).

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Craig Topper (topperc)

Changes

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).

---

Full diff: https://github.com/llvm/llvm-project/pull/98533.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp (+9-4)
• (modified) llvm/test/Transforms/InstCombine/add.ll (+15)

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 8b3ecaf25cd9d..adb61cd4b9692 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:    [[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
+  %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) {

[dtcxzyw]

LGTM. Thank you!

[nikic]

Sounds reasonable.