[InstCombine] Preserve flags for abs(X) * abs(X) and nabs(X) * nabs(X)

Alive2 Proofs:

https://alive2.llvm.org/ce/z/CN9BP-
https://alive2.llvm.org/ce/z/kq9vh6
https://alive2.llvm.org/ce/z/LtVSL5

Author: AZero13

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

@@ -320,8 +320,26 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
 
   // abs(X) * abs(X) -> X * X
   Value *X;
-  if (Op0 == Op1 && match(Op0, m_Intrinsic<Intrinsic::abs>(m_Value(X))))
-    return BinaryOperator::CreateMul(X, X);
+  if (Op0 == Op1) {
+    // nabs(X) * nabs(X) -> X * X
+    if (match(Op0, m_Neg(m_Intrinsic<Intrinsic::abs>(m_Value(X))))) {
+      Instruction *NewMul = BinaryOperator::CreateMul(X, X);
+      if (I.hasNoUnsignedWrap()) {
+        // (nabs * nabs) nuw in this case also implies nsw
+        NewMul->setHasNoUnsignedWrap(true);
+        NewMul->setHasNoSignedWrap(true);
+      } else if (I.hasNoSignedWrap())
+        NewMul->setHasNoSignedWrap(true);
+
+      return NewMul;
+    }
+
+    if (match(Op0, m_Intrinsic<Intrinsic::abs>(m_Value(X)))) {
+      Instruction *NewMul = BinaryOperator::CreateMul(X, X);
+      NewMul->setHasNoSignedWrap(I.hasNoSignedWrap());
+      return NewMul;
+    }
+  }
 
   {
     Value *Y;

llvm/test/Transforms/InstCombine/mul-inseltpoison.ll

@@ -927,7 +927,7 @@ define i32 @combine_mul_abs_intrin_nuw(i32 %x) {
 
 define i32 @combine_mul_abs_intrin_nsw(i32 %x) {
 ; CHECK-LABEL: @combine_mul_abs_intrin_nsw(
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
@@ -948,7 +948,7 @@ define i32 @combine_mul_nabs_intrin(i32 %x) {
 
 define i32 @combine_mul_nabs_intrin_nuw(i32 %x) {
 ; CHECK-LABEL: @combine_mul_nabs_intrin_nuw(
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
@@ -959,7 +959,7 @@ define i32 @combine_mul_nabs_intrin_nuw(i32 %x) {
 
 define i32 @combine_mul_nabs_intrin_nsw(i32 %x) {
 ; CHECK-LABEL: @combine_mul_nabs_intrin_nsw(
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)

llvm/test/Transforms/InstCombine/mul.ll

@@ -1607,7 +1607,7 @@ define i32 @combine_mul_nabs_i32(i32 %0) {
 
 define i32 @combine_mul_nabs_i32_unsigned_wrap(i32 %0) {
 ; CHECK-LABEL: @combine_mul_nabs_i32_unsigned_wrap(
-; CHECK-NEXT:    [[M:%.*]] = mul i32 [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    [[M:%.*]] = mul nuw nsw i32 [[TMP0:%.*]], [[TMP0]]
 ; CHECK-NEXT:    ret i32 [[M]]
 ;
   %c = icmp slt i32 %0, 0
@@ -1619,7 +1619,7 @@ define i32 @combine_mul_nabs_i32_unsigned_wrap(i32 %0) {
 
 define i32 @combine_mul_nabs_i32_signed_wrap(i32 %0) {
 ; CHECK-LABEL: @combine_mul_nabs_i32_signed_wrap(
-; CHECK-NEXT:    [[M:%.*]] = mul i32 [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    [[M:%.*]] = mul nsw i32 [[TMP0:%.*]], [[TMP0]]
 ; CHECK-NEXT:    ret i32 [[M]]
 ;
   %c = icmp slt i32 %0, 0
@@ -1644,7 +1644,7 @@ define <4 x i32> @combine_mul_nabs_v4i32(<4 x i32> %0) {
 
 define <4 x i32> @combine_mul_nabs_v4i32_unsigned_wrap(<4 x i32> %0) {
 ; CHECK-LABEL: @combine_mul_nabs_v4i32_unsigned_wrap(
-; CHECK-NEXT:    [[M:%.*]] = mul <4 x i32> [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    [[M:%.*]] = mul nuw nsw <4 x i32> [[TMP0:%.*]], [[TMP0]]
 ; CHECK-NEXT:    ret <4 x i32> [[M]]
 ;
   %c = icmp slt <4 x i32> %0, zeroinitializer
@@ -1656,7 +1656,7 @@ define <4 x i32> @combine_mul_nabs_v4i32_unsigned_wrap(<4 x i32> %0) {
 
 define <4 x i32> @combine_mul_nabs_v4i32_signed_wrap(<4 x i32> %0) {
 ; CHECK-LABEL: @combine_mul_nabs_v4i32_signed_wrap(
-; CHECK-NEXT:    [[M:%.*]] = mul <4 x i32> [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    [[M:%.*]] = mul nsw <4 x i32> [[TMP0:%.*]], [[TMP0]]
 ; CHECK-NEXT:    ret <4 x i32> [[M]]
 ;
   %c = icmp slt <4 x i32> %0, zeroinitializer
@@ -1689,7 +1689,7 @@ define i32 @combine_mul_nabs_intrin(i32 %x) {
 
 define i32 @combine_mul_nabs_intrin_unsigned_flags(i32 %x) {
 ; CHECK-LABEL: @combine_mul_nabs_intrin_unsigned_flags(
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
@@ -1700,7 +1700,7 @@ define i32 @combine_mul_nabs_intrin_unsigned_flags(i32 %x) {
 
 define i32 @combine_mul_nabs_intrin_signed_flags(i32 %x) {
 ; CHECK-LABEL: @combine_mul_nabs_intrin_signed_flags(
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)