[llvm] [InstCombine] abs(X) * abs(X) and nabs(X) * nabs(X) should preserve flags (PR #88662)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: AtariDreams (AtariDreams)

<details>
<summary>Changes</summary>

Alive2 Proofs:

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

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp (+13-4) 
- (modified) llvm/test/Transforms/InstCombine/mul.ll (+72-1) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 4dc1319f1c437f..ee092add097771 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -323,11 +323,20 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   if (Op0 == Op1) {
     Value *X, *Y;
     SelectPatternFlavor SPF = matchSelectPattern(Op0, X, Y).Flavor;
-    if (SPF == SPF_ABS || SPF == SPF_NABS)
-      return BinaryOperator::CreateMul(X, X);
+    if (SPF == SPF_NABS ||
+        match(Op0, m_Neg(m_Intrinsic<Intrinsic::abs>(m_Value(X))))) {
+      if (I.hasNoUnsignedWrap()) {
+        Instruction *NewMul = BinaryOperator::CreateNUWMul(X, X);
+        NewMul->setHasNoSignedWrap(true);
+        return NewMul;
+      }
+      return I.hasNoSignedWrap() ? BinaryOperator::CreateNSWMul(X, X)
+                                 : BinaryOperator::CreateMul(X, X);
+    }
 
-    if (match(Op0, m_Intrinsic<Intrinsic::abs>(m_Value(X))))
-      return BinaryOperator::CreateMul(X, X);
+    if (SPF == SPF_ABS || match(Op0, m_Intrinsic<Intrinsic::abs>(m_Value(X))))
+      return I.hasNoSignedWrap() ? BinaryOperator::CreateNSWMul(X, X)
+                                 : BinaryOperator::CreateMul(X, X);
   }
 
   {
diff --git a/llvm/test/Transforms/InstCombine/mul.ll b/llvm/test/Transforms/InstCombine/mul.ll
index d4a689c60786e0..e478e34ddd4859 100644
--- a/llvm/test/Transforms/InstCombine/mul.ll
+++ b/llvm/test/Transforms/InstCombine/mul.ll
@@ -1605,6 +1605,31 @@ define i32 @combine_mul_nabs_i32(i32 %0) {
   ret i32 %m
 }
 
+define i32 @combine_mul_nabs_i32_unsigned_wrap(i32 %0) {
+; CHECK-LABEL: @combine_mul_nabs_i32_unsigned_wrap(
+; CHECK-NEXT:    [[M:%.*]] = mul nuw nsw i32 [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret i32 [[M]]
+;
+  %c = icmp slt i32 %0, 0
+  %s = sub nsw i32 0, %0
+  %r = select i1 %c, i32 %0, i32 %s
+  %m = mul nuw i32 %r, %r
+  ret i32 %m
+}
+
+define i32 @combine_mul_nabs_i32_signed_wrap(i32 %0) {
+; CHECK-LABEL: @combine_mul_nabs_i32_signed_wrap(
+; CHECK-NEXT:    [[M:%.*]] = mul nsw i32 [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret i32 [[M]]
+;
+  %c = icmp slt i32 %0, 0
+  %s = sub nsw i32 0, %0
+  %r = select i1 %c, i32 %0, i32 %s
+  %m = mul nsw i32 %r, %r
+  ret i32 %m
+}
+
+
 define <4 x i32> @combine_mul_nabs_v4i32(<4 x i32> %0) {
 ; CHECK-LABEL: @combine_mul_nabs_v4i32(
 ; CHECK-NEXT:    [[M:%.*]] = mul <4 x i32> [[TMP0:%.*]], [[TMP0]]
@@ -1617,13 +1642,37 @@ define <4 x i32> @combine_mul_nabs_v4i32(<4 x i32> %0) {
   ret <4 x i32> %m
 }
 
+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 nuw nsw <4 x i32> [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret <4 x i32> [[M]]
+;
+  %c = icmp slt <4 x i32> %0, zeroinitializer
+  %s = sub nsw <4 x i32> zeroinitializer, %0
+  %r = select <4 x i1> %c, <4 x i32> %0, <4 x i32> %s
+  %m = mul nuw <4 x i32> %r, %r
+  ret <4 x i32> %m
+}
+
+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 nsw <4 x i32> [[TMP0:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret <4 x i32> [[M]]
+;
+  %c = icmp slt <4 x i32> %0, zeroinitializer
+  %s = sub nsw <4 x i32> zeroinitializer, %0
+  %r = select <4 x i1> %c, <4 x i32> %0, <4 x i32> %s
+  %m = mul nsw <4 x i32> %r, %r
+  ret <4 x i32> %m
+}
+
 define i32 @combine_mul_abs_intrin(i32 %x) {
 ; CHECK-LABEL: @combine_mul_abs_intrin(
 ; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[X:%.*]], [[X]]
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
-  %mul = mul i32 %abs, %abs
+  %mul = mul nuw i32 %abs, %abs
   ret i32 %mul
 }
 
@@ -1638,6 +1687,28 @@ define i32 @combine_mul_nabs_intrin(i32 %x) {
   ret i32 %mul
 }
 
+define i32 @combine_mul_nabs_intrin_unsigned_flags(i32 %x) {
+; CHECK-LABEL: @combine_mul_nabs_intrin_unsigned_flags(
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    ret i32 [[MUL]]
+;
+  %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
+  %neg = sub i32 0, %abs
+  %mul = mul nuw i32 %neg, %neg
+  ret i32 %mul
+}
+
+define i32 @combine_mul_nabs_intrin_signed_flags(i32 %x) {
+; CHECK-LABEL: @combine_mul_nabs_intrin_signed_flags(
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[X:%.*]], [[X]]
+; CHECK-NEXT:    ret i32 [[MUL]]
+;
+  %abs = call i32 @llvm.abs.i32(i32 %x, i1 false)
+  %neg = sub i32 0, %abs
+  %mul = mul nsw i32 %neg, %neg
+  ret i32 %mul
+}
+
 ; z * splat(0) = splat(0), even for scalable vectors
 define <vscale x 2 x i64> @mul_scalable_splat_zero(<vscale x 2 x i64> %z) {
 ; CHECK-LABEL: @mul_scalable_splat_zero(

``````````

</details>


https://github.com/llvm/llvm-project/pull/88662