[llvm] [InstCombine] Propagate NSW/NUW flags for `(X - Y) - Z -> X - (Y + Z)` (PR #72693)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/72693

Alive2: https://alive2.llvm.org/ce/z/gqeaVo

Related patch: https://github.com/llvm/llvm-project/commit/31d219d2997fed1b7dc97e0adf170d5aaf65883e

>From 8ff3309ab8950c467833520564d76f76779b6d06 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 18 Nov 2023 02:28:23 +0800
Subject: [PATCH 1/2] [InstCombine] Add extra pre-commit tests. NFC.

---
 .../Transforms/InstCombine/sub-from-sub.ll    | 26 ++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/InstCombine/sub-from-sub.ll b/llvm/test/Transforms/InstCombine/sub-from-sub.ll
index 7976cdc970b9c5b..a045ebfd45d6f88 100644
--- a/llvm/test/Transforms/InstCombine/sub-from-sub.ll
+++ b/llvm/test/Transforms/InstCombine/sub-from-sub.ll
@@ -17,7 +17,7 @@ define i8 @t0(i8 %x, i8 %y, i8 %z) {
   ret i8 %r
 }
 
-; No flags are propagated
+; NSW/NUW flags are propagated
 define i8 @t1_flags(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @t1_flags(
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], [[Z:%.*]]
@@ -29,6 +29,30 @@ define i8 @t1_flags(i8 %x, i8 %y, i8 %z) {
   ret i8 %r
 }
 
+; NUW flags are propagated
+define i8 @t1_flags_nuw_only(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @t1_flags_nuw_only(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = sub i8 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %o0 = sub nuw i8 %x, %y
+  %r = sub nuw i8 %o0, %z
+  ret i8 %r
+}
+
+; Negative test: NSW flags cannot be propagated
+define i8 @t1_flags_nsw_only(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @t1_flags_nsw_only(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = sub i8 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %o0 = sub nsw i8 %x, %y
+  %r = sub nsw i8 %o0, %z
+  ret i8 %r
+}
+
 ; The inner sub must have single use.
 define i8 @n2(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @n2(

>From 64a17ea0f45b856ff995256ade5a331a580be6b4 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 18 Nov 2023 03:29:01 +0800
Subject: [PATCH 2/2] [InstCombine] Propagate NSW/NUW flags for `(X - Y) - Z ->
 X - (Y + Z)`

---
 llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp | 11 +++++++++--
 llvm/test/Transforms/InstCombine/sub-from-sub.ll      |  8 ++++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index 318992b55e4f9f8..90b1c133461a408 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -2191,8 +2191,15 @@ Instruction *InstCombinerImpl::visitSub(BinaryOperator &I) {
 
   // ((X - Y) - Op1)  -->  X - (Y + Op1)
   if (match(Op0, m_OneUse(m_Sub(m_Value(X), m_Value(Y))))) {
-    Value *Add = Builder.CreateAdd(Y, Op1);
-    return BinaryOperator::CreateSub(X, Add);
+    OverflowingBinaryOperator *LHSSub = cast<OverflowingBinaryOperator>(Op0);
+    bool HasNUW = I.hasNoUnsignedWrap() && LHSSub->hasNoUnsignedWrap();
+    bool HasNSW = HasNUW && I.hasNoSignedWrap() && LHSSub->hasNoSignedWrap();
+    Value *Add = Builder.CreateAdd(Y, Op1, "", /* HasNUW */ HasNUW,
+                                   /* HasNSW */ HasNSW);
+    BinaryOperator *Sub = BinaryOperator::CreateSub(X, Add);
+    Sub->setHasNoUnsignedWrap(HasNUW);
+    Sub->setHasNoSignedWrap(HasNSW);
+    return Sub;
   }
 
   // (~X) - (~Y) --> Y - X
diff --git a/llvm/test/Transforms/InstCombine/sub-from-sub.ll b/llvm/test/Transforms/InstCombine/sub-from-sub.ll
index a045ebfd45d6f88..b5e978cd86cbe15 100644
--- a/llvm/test/Transforms/InstCombine/sub-from-sub.ll
+++ b/llvm/test/Transforms/InstCombine/sub-from-sub.ll
@@ -20,8 +20,8 @@ define i8 @t0(i8 %x, i8 %y, i8 %z) {
 ; NSW/NUW flags are propagated
 define i8 @t1_flags(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @t1_flags(
-; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = sub i8 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw nsw i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = sub nuw nsw i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %o0 = sub nuw nsw i8 %x, %y
@@ -32,8 +32,8 @@ define i8 @t1_flags(i8 %x, i8 %y, i8 %z) {
 ; NUW flags are propagated
 define i8 @t1_flags_nuw_only(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @t1_flags_nuw_only(
-; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = sub i8 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = sub nuw i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %o0 = sub nuw i8 %x, %y