[InstCombine] extractvalue (any_mul_with_overflow X, 2^n), 0 -> X << n

[Chenbing.Zheng]

Summary

Alive2: https://alive2.llvm.org/ce/z/JLmabt (umul)

https://alive2.llvm.org/ce/z/J_ruXR  (smul)
https://alive2.llvm.org/ce/z/o9SVSz (vector)

[spatel]

Why not do the transform with smul too?

[RKSimon]

Is this a better proof? https://alive2.llvm.org/ce/z/JLmabt (noundefs) https://alive2.llvm.org/ce/z/J_ruXR (with undefs)

[Chenbing.Zheng]

add smul support

[Chenbing.Zheng]

In D133188#3767207, @RKSimon wrote:

Is this a better proof? https://alive2.llvm.org/ce/z/JLmabt (noundefs) https://alive2.llvm.org/ce/z/J_ruXR (with undefs)

Yer, thanks. I have update it to summary.

[spatel]

llvm/test/Transforms/InstCombine/with_overflow.ll

1007

Wrong comment or wrong test?

[Chenbing.Zheng]

fix comment

[spatel]

This change allows vectors with undef elements, but there are no tests for that. Please adjust at least one test to exercise that pattern - a vector type with constant that includes a poison element.
You can also verify a specific constant with poison with Alive2 and add a link to the patch summary.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3285–3286

m_AllOnes ?

[Chenbing.Zheng]

add vector tests

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3285–3286

We don't just deal with AllOnes

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

 }
 
 Instruction *
 InstCombinerImpl::foldExtractOfOverflowIntrinsic(ExtractValueInst &EV) {
   auto *WO = dyn_cast<WithOverflowInst>(EV.getAggregateOperand());
   if (!WO)
     return nullptr;
 
-  // extractvalue (any_mul_with_overflow X, -1), 0 --> -X
   Intrinsic::ID OvID = WO->getIntrinsicID();
   const APInt *C = nullptr;
   if (match(WO->getRHS(), m_APIntAllowUndef(C))) {
-    if (*EV.idx_begin() == 0 &&
-        (OvID == Intrinsic::smul_with_overflow ||
-         OvID == Intrinsic::umul_with_overflow) &&
-        C->isAllOnes()) {
-      return BinaryOperator::CreateNeg(WO->getArgOperand(0));
+    if (*EV.idx_begin() == 0 && (OvID == Intrinsic::smul_with_overflow ||
+                                 OvID == Intrinsic::umul_with_overflow)) {
+      // extractvalue (any_mul_with_overflow X, -1), 0 --> -X
+      if (C->isAllOnes())
+        return BinaryOperator::CreateNeg(WO->getLHS());
+      // extractvalue (any_mul_with_overflow X, 2^n), 0 --> X << n
+      if (C->isPowerOf2()) {
+        return BinaryOperator::CreateShl(
+            WO->getLHS(),
+            ConstantInt::get(WO->getLHS()->getType(), C->logBase2()));
+      }
     }
   }
 
   // We're extracting from an overflow intrinsic. See if we're the only user.
   // That allows us to simplify multiple result intrinsics to simpler things
   // that just get one value.
   if (!WO->hasOneUse())
     return nullptr;

llvm/test/Transforms/InstCombine/with_overflow.ll

;
   %ov = extractvalue { i32, i1 } %mul, 1
   %res = extractvalue { i32, i1 } %mul, 0
   %ret = select i1 %ov, i32 -1, i32 %res
   ret i32 %ret
 }
 
 define i32 @umultest4(i32 %n) nounwind {
 ; CHECK-LABEL: @umultest4(
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[N:%.*]], 1
-; CHECK-NEXT:    [[MUL:%.*]] = call { i32, i1 } @llvm.umul.with.overflow.i32(i32 [[SHR]], i32 4)
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i32, i1 } [[MUL]], 1
-; CHECK-NEXT:    [[RES:%.*]] = extractvalue { i32, i1 } [[MUL]], 0
-; CHECK-NEXT:    [[RET:%.*]] = select i1 [[OV]], i32 -1, i32 [[RES]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[N:%.*]], 1
+; CHECK-NEXT:    [[RES:%.*]] = and i32 [[TMP1]], -4
+; CHECK-NEXT:    [[OV_INV:%.*]] = icmp sgt i32 [[N]], -1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[OV_INV]], i32 [[RES]], i32 -1
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %shr = lshr i32 %n, 1
   %mul = call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %shr, i32 4)
   %ov = extractvalue { i32, i1 } %mul, 1
   %res = extractvalue { i32, i1 } %mul, 0
   %ret = select i1 %ov, i32 -1, i32 %res
   ret i32 %ret

 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 -1)
   %m0 = extractvalue { i8, i1 } %m, 0
   %m1 = extractvalue { i8, i1 } %m, 1
   %r = select i1 %m1, i8 -1, i8 %m0
   ret i8 %r
 }
 
-; TODO: this could be 'shl' and 'icmp'
-
 define i8 @umul_2(i8 %x, i1* %p) {
 ; CHECK-LABEL: @umul_2(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 [[X:%.*]], i8 2)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 1
+; CHECK-NEXT:    [[OV:%.*]] = icmp slt i8 [[X]], 0
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 2)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @umul_8(i8 %x, i1* %p) {
 ; CHECK-LABEL: @umul_8(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 [[X:%.*]], i8 8)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 3
+; CHECK-NEXT:    [[OV:%.*]] = icmp ugt i8 [[X]], 31
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 8)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @umul_64(i8 %x, i1* %p) {
 ; CHECK-LABEL: @umul_64(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 [[X:%.*]], i8 64)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 6
+; CHECK-NEXT:    [[OV:%.*]] = icmp ugt i8 [[X]], 3
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 64)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @umul_256(i8 %x, i1* %p) {
 ; CHECK-LABEL: @umul_256(
 ; CHECK-NEXT:    store i1 false, i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 0
 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 256)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
+define <4 x i8> @umul_4_vec_poison(<4 x i8> %x, <4 x i1>* %p) {
+; CHECK-LABEL: @umul_4_vec_poison(
+; CHECK-NEXT:    [[R:%.*]] = shl <4 x i8> [[X:%.*]], <i8 2, i8 2, i8 2, i8 2>
+; CHECK-NEXT:    [[OV:%.*]] = icmp ugt <4 x i8> [[X]], <i8 63, i8 63, i8 63, i8 63>
+; CHECK-NEXT:    store <4 x i1> [[OV]], <4 x i1>* [[P:%.*]], align 1
+; CHECK-NEXT:    ret <4 x i8> [[R]]
+;
+  %m = tail call { <4 x i8>, <4 x i1> } @llvm.umul.with.overflow.v4i8(<4 x i8> %x, <4 x i8> <i8 poison, i8 4, i8 4, i8 poison>)
+  %r = extractvalue { <4 x i8>, <4 x i1> } %m, 0
+  %ov = extractvalue { <4 x i8>, <4 x i1> } %m, 1
+  store <4 x i1> %ov, <4 x i1>* %p
+  ret <4 x i8> %r
+}
+
+; Negative test: not PowerOf2
+
 define i8 @umul_3(i8 %x, i1* %p) {
 ; CHECK-LABEL: @umul_3(
 ; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 [[X:%.*]], i8 3)
 ; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
 ; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.umul.with.overflow.i8(i8 %x, i8 3)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @smul_4(i8 %x, i1* %p) {
 ; CHECK-LABEL: @smul_4(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 [[X:%.*]], i8 4)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], -32
+; CHECK-NEXT:    [[OV:%.*]] = icmp ult i8 [[TMP1]], -64
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 %x, i8 4)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @smul_16(i8 %x, i1* %p) {
 ; CHECK-LABEL: @smul_16(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 [[X:%.*]], i8 16)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 4
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], -8
+; CHECK-NEXT:    [[OV:%.*]] = icmp ult i8 [[TMP1]], -16
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 %x, i8 16)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @smul_32(i8 %x, i1* %p) {
 ; CHECK-LABEL: @smul_32(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 [[X:%.*]], i8 32)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 5
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], -4
+; CHECK-NEXT:    [[OV:%.*]] = icmp ult i8 [[TMP1]], -8
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 %x, i8 32)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
 define i8 @smul_128(i8 %x, i1* %p) {
 ; CHECK-LABEL: @smul_128(
-; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 [[X:%.*]], i8 -128)
-; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
-; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], 7
+; CHECK-NEXT:    [[OV:%.*]] = icmp ugt i8 [[X]], 1
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 %x, i8 128)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }
 
+define <4 x i8> @smul_2_vec_poison(<4 x i8> %x, <4 x i1>* %p) {
+; CHECK-LABEL: @smul_2_vec_poison(
+; CHECK-NEXT:    [[R:%.*]] = shl <4 x i8> [[X:%.*]], <i8 1, i8 1, i8 1, i8 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = add <4 x i8> [[X]], <i8 64, i8 64, i8 64, i8 64>
+; CHECK-NEXT:    [[OV:%.*]] = icmp slt <4 x i8> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    store <4 x i1> [[OV]], <4 x i1>* [[P:%.*]], align 1
+; CHECK-NEXT:    ret <4 x i8> [[R]]
+;
+  %m = tail call { <4 x i8>, <4 x i1> } @llvm.smul.with.overflow.v4i8(<4 x i8> %x, <4 x i8> <i8 poison, i8 2, i8 2, i8 poison>)
+  %r = extractvalue { <4 x i8>, <4 x i1> } %m, 0
+  %ov = extractvalue { <4 x i8>, <4 x i1> } %m, 1
+  store <4 x i1> %ov, <4 x i1>* %p
+  ret <4 x i8> %r
+}
+
+; Negative test: not PowerOf2
+
 define i8 @smul_7(i8 %x, i1* %p) {
 ; CHECK-LABEL: @smul_7(
 ; CHECK-NEXT:    [[M:%.*]] = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 [[X:%.*]], i8 7)
 ; CHECK-NEXT:    [[R:%.*]] = extractvalue { i8, i1 } [[M]], 0
 ; CHECK-NEXT:    [[OV:%.*]] = extractvalue { i8, i1 } [[M]], 1
 ; CHECK-NEXT:    store i1 [[OV]], i1* [[P:%.*]], align 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = tail call { i8, i1 } @llvm.smul.with.overflow.i8(i8 %x, i8 7)
   %r = extractvalue { i8, i1 } %m, 0
   %ov = extractvalue { i8, i1 } %m, 1
   store i1 %ov, i1* %p
   ret i8 %r
 }