From 8146d2023c9281dea9810212135283710e53451e Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Mon, 24 Jun 2019 21:37:03 +0000
Subject: [PATCH] InstCombine: Preserve nuw when reassociating nuw ops [3/3]

Alive says this is OK.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@364235 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../InstCombine/InstructionCombining.cpp      | 42 +++++++++------
 .../Transforms/InstCombine/reassociate-nuw.ll | 52 ++++++++++++++++++-
 2 files changed, 77 insertions(+), 17 deletions(-)

diff --git a/lib/Transforms/InstCombine/InstructionCombining.cpp b/lib/Transforms/InstCombine/InstructionCombining.cpp
index fcbe738aad7..385f4926b84 100644
--- a/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -590,32 +590,44 @@ Value *InstCombiner::tryFactorization(BinaryOperator &I,
     ++NumFactor;
     SimplifiedInst->takeName(&I);
 
-    // Check if we can add NSW flag to SimplifiedInst. If so, set NSW flag.
-    // TODO: Check for NUW.
+    // Check if we can add NSW/NUW flags to SimplifiedInst. If so, set them.
     if (BinaryOperator *BO = dyn_cast<BinaryOperator>(SimplifiedInst)) {
       if (isa<OverflowingBinaryOperator>(SimplifiedInst)) {
         bool HasNSW = false;
-        if (isa<OverflowingBinaryOperator>(&I))
+        bool HasNUW = false;
+        if (isa<OverflowingBinaryOperator>(&I)) {
           HasNSW = I.hasNoSignedWrap();
+          HasNUW = I.hasNoUnsignedWrap();
+        }
 
-        if (auto *LOBO = dyn_cast<OverflowingBinaryOperator>(LHS))
+        if (auto *LOBO = dyn_cast<OverflowingBinaryOperator>(LHS)) {
           HasNSW &= LOBO->hasNoSignedWrap();
+          HasNUW &= LOBO->hasNoUnsignedWrap();
+        }
 
-        if (auto *ROBO = dyn_cast<OverflowingBinaryOperator>(RHS))
+        if (auto *ROBO = dyn_cast<OverflowingBinaryOperator>(RHS)) {
           HasNSW &= ROBO->hasNoSignedWrap();
+          HasNUW &= ROBO->hasNoUnsignedWrap();
+        }
 
-        // We can propagate 'nsw' if we know that
-        //  %Y = mul nsw i16 %X, C
-        //  %Z = add nsw i16 %Y, %X
-        // =>
-        //  %Z = mul nsw i16 %X, C+1
-        //
-        // iff C+1 isn't INT_MIN
         const APInt *CInt;
         if (TopLevelOpcode == Instruction::Add &&
-            InnerOpcode == Instruction::Mul)
-          if (match(V, m_APInt(CInt)) && !CInt->isMinSignedValue())
-            BO->setHasNoSignedWrap(HasNSW);
+            InnerOpcode == Instruction::Mul) {
+          // We can propagate 'nsw' if we know that
+          //  %Y = mul nsw i16 %X, C
+          //  %Z = add nsw i16 %Y, %X
+          // =>
+          //  %Z = mul nsw i16 %X, C+1
+          //
+          // iff C+1 isn't INT_MIN
+          if (match(V, m_APInt(CInt))) {
+            if (!CInt->isMinSignedValue())
+              BO->setHasNoSignedWrap(HasNSW);
+          }
+
+          // nuw can be propagated with any constant or nuw value.
+          BO->setHasNoUnsignedWrap(HasNUW);
+        }
       }
     }
   }
diff --git a/test/Transforms/InstCombine/reassociate-nuw.ll b/test/Transforms/InstCombine/reassociate-nuw.ll
index 22ad8c259d4..a40d427fd29 100644
--- a/test/Transforms/InstCombine/reassociate-nuw.ll
+++ b/test/Transforms/InstCombine/reassociate-nuw.ll
@@ -92,7 +92,7 @@ define i32 @reassoc_x2_sub_nuw(i32 %x, i32 %y) {
 
 define i32 @tryFactorization_add_nuw_mul_nuw(i32 %x) {
 ; CHECK-LABEL: @tryFactorization_add_nuw_mul_nuw(
-; CHECK-NEXT:    [[ADD2:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[ADD2:%.*]] = shl nuw i32 [[X:%.*]], 2
 ; CHECK-NEXT:    ret i32 [[ADD2]]
 ;
   %mul1 = mul nuw i32 %x, 3
@@ -102,7 +102,7 @@ define i32 @tryFactorization_add_nuw_mul_nuw(i32 %x) {
 
 define i32 @tryFactorization_add_nuw_mul_nuw_int_max(i32 %x) {
 ; CHECK-LABEL: @tryFactorization_add_nuw_mul_nuw_int_max(
-; CHECK-NEXT:    [[ADD2:%.*]] = shl i32 [[X:%.*]], 31
+; CHECK-NEXT:    [[ADD2:%.*]] = shl nuw i32 [[X:%.*]], 31
 ; CHECK-NEXT:    ret i32 [[ADD2]]
 ;
   %mul1 = mul nuw i32 %x, 2147483647
@@ -129,3 +129,51 @@ define i32 @tryFactorization_add_nuw_mul(i32 %x) {
   %add2 = add i32 %mul1, %x
   ret i32 %add2
 }
+
+define i32 @tryFactorization_add_nuw_mul_nuw_mul_nuw_var(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: @tryFactorization_add_nuw_mul_nuw_mul_nuw_var(
+; CHECK-NEXT:    [[MUL21:%.*]] = add i32 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = mul nuw i32 [[MUL21]], [[X:%.*]]
+; CHECK-NEXT:    ret i32 [[ADD1]]
+;
+  %mul1 = mul nuw i32 %x, %y
+  %mul2 = mul nuw i32 %x, %z
+  %add1 = add nuw i32 %mul1, %mul2
+  ret i32 %add1
+}
+
+define i32 @tryFactorization_add_nuw_mul_mul_nuw_var(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: @tryFactorization_add_nuw_mul_mul_nuw_var(
+; CHECK-NEXT:    [[MUL21:%.*]] = add i32 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = mul i32 [[MUL21]], [[X:%.*]]
+; CHECK-NEXT:    ret i32 [[ADD1]]
+;
+  %mul1 = mul i32 %x, %y
+  %mul2 = mul nuw i32 %x, %z
+  %add1 = add nuw i32 %mul1, %mul2
+  ret i32 %add1
+}
+
+define i32 @tryFactorization_add_nuw_mul_nuw_mul_var(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: @tryFactorization_add_nuw_mul_nuw_mul_var(
+; CHECK-NEXT:    [[MUL21:%.*]] = add i32 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = mul i32 [[MUL21]], [[X:%.*]]
+; CHECK-NEXT:    ret i32 [[ADD1]]
+;
+  %mul1 = mul nuw i32 %x, %y
+  %mul2 = mul i32 %x, %z
+  %add1 = add nuw i32 %mul1, %mul2
+  ret i32 %add1
+}
+
+define i32 @tryFactorization_add_mul_nuw_mul_var(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: @tryFactorization_add_mul_nuw_mul_var(
+; CHECK-NEXT:    [[MUL21:%.*]] = add i32 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = mul i32 [[MUL21]], [[X:%.*]]
+; CHECK-NEXT:    ret i32 [[ADD1]]
+;
+  %mul1 = mul nuw i32 %x, %y
+  %mul2 = mul nuw i32 %x, %z
+  %add1 = add i32 %mul1, %mul2
+  ret i32 %add1
+}
-- 
2.40.0