[InstCombine] Support (mul (sext x), cst) --> (sext (mul x, cst')) and (mul (zext x), cst) --> (zext (mul x, cst')) for vectors constants.

[craig.topper]

Summary

Similar to D51236, but for mul instead of add.

[spatel]

LGTM

[spatel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

328–332

Nit: here and below, it would probably be easier to read and shorter using 'match' rather than dyn_cast.

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Instruction *InstCombiner::visitMul(BinaryOperator &I) {
   const APInt *C;
   if (match(Op0, m_LShr(m_Value(X), m_APInt(C))) && *C == C->getBitWidth() - 1)
     return BinaryOperator::CreateAnd(Builder.CreateAShr(X, *C), Op1);
   if (match(Op1, m_LShr(m_Value(X), m_APInt(C))) && *C == C->getBitWidth() - 1)
     return BinaryOperator::CreateAnd(Builder.CreateAShr(X, *C), Op0);
 
   // Check for (mul (sext x), y), see if we can merge this into an
   // integer mul followed by a sext.
   if (SExtInst *Op0Conv = dyn_cast<SExtInst>(Op0)) {
     // (mul (sext x), cst) --> (sext (mul x, cst'))
-    if (ConstantInt *Op1C = dyn_cast<ConstantInt>(Op1)) {
+    Constant *Op1C;
+    if (match(Op1, m_Constant(Op1C))) {
       if (Op0Conv->hasOneUse()) {
         Constant *CI =
             ConstantExpr::getTrunc(Op1C, Op0Conv->getOperand(0)->getType());
         if (ConstantExpr::getSExt(CI, I.getType()) == Op1C &&
             willNotOverflowSignedMul(Op0Conv->getOperand(0), CI, I)) {
           // Insert the new, smaller mul.
           Value *NewMul =
               Builder.CreateNSWMul(Op0Conv->getOperand(0), CI, "mulconv");
           return new SExtInst(NewMul, I.getType());
if (SExtInst *Op1Conv = dyn_cast<SExtInst>(Op1)) {
       }
     }
   }
 
   // Check for (mul (zext x), y), see if we can merge this into an
   // integer mul followed by a zext.
   if (auto *Op0Conv = dyn_cast<ZExtInst>(Op0)) {
     // (mul (zext x), cst) --> (zext (mul x, cst'))
-    if (ConstantInt *Op1C = dyn_cast<ConstantInt>(Op1)) {
+    Constant *Op1C;
+    if (match(Op1, m_Constant(Op1C))) {
       if (Op0Conv->hasOneUse()) {
         Constant *CI =
             ConstantExpr::getTrunc(Op1C, Op0Conv->getOperand(0)->getType());
         if (ConstantExpr::getZExt(CI, I.getType()) == Op1C &&
             willNotOverflowUnsignedMul(Op0Conv->getOperand(0), CI, I)) {
           // Insert the new, smaller mul.
           Value *NewMul =
               Builder.CreateNUWMul(Op0Conv->getOperand(0), CI, "mulconv");

test/Transforms/InstCombine/sink-zext.ll

;
   %sext = sext i32 %ashr to i64
   %mul = mul i64 %sext, 32767
   ret i64 %mul
 }
 
 define <2 x i64> @test8_splat(<2 x i32> %V) {
 ; CHECK-LABEL: @test8_splat(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i32> [[ASHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw <2 x i64> [[SEXT]], <i64 32767, i64 32767>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw <2 x i32> [[ASHR]], <i32 32767, i32 32767>
+; CHECK-NEXT:    [[MUL:%.*]] = sext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %ashr = ashr <2 x i32> %V, <i32 16, i32 16>
   %sext = sext <2 x i32> %ashr to <2 x i64>
   %mul = mul <2 x i64> %sext, <i64 32767, i64 32767>
   ret <2 x i64> %mul
 }
 
 define <2 x i64> @test8_vec(<2 x i32> %V) {
 ; CHECK-LABEL: @test8_vec(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i32> [[ASHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw <2 x i64> [[SEXT]], <i64 32767, i64 16384>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw <2 x i32> [[ASHR]], <i32 32767, i32 16384>
+; CHECK-NEXT:    [[MUL:%.*]] = sext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %ashr = ashr <2 x i32> %V, <i32 16, i32 16>
   %sext = sext <2 x i32> %ashr to <2 x i64>
   %mul = mul <2 x i64> %sext, <i64 32767, i64 16384>
   ret <2 x i64> %mul
 }
 
 define <2 x i64> @test8_vec2(<2 x i32> %V) {
 ; CHECK-LABEL: @test8_vec2(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i32> [[ASHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw <2 x i64> [[SEXT]], <i64 32767, i64 -32767>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw <2 x i32> [[ASHR]], <i32 32767, i32 -32767>
+; CHECK-NEXT:    [[MUL:%.*]] = sext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %ashr = ashr <2 x i32> %V, <i32 16, i32 16>
   %sext = sext <2 x i32> %ashr to <2 x i64>
   %mul = mul <2 x i64> %sext, <i64 32767, i64 -32767>
   ret <2 x i64> %mul
 }
 
 define i64 @test9(i32 %V) {
 ; CHECK-LABEL: @test9(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr i32 [[V:%.*]], 16
 ; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw i32 [[ASHR]], -32767
 ; CHECK-NEXT:    [[MUL:%.*]] = sext i32 [[MULCONV]] to i64
 ; CHECK-NEXT:    ret i64 [[MUL]]
 ;
   %ashr = ashr i32 %V, 16
   %sext = sext i32 %ashr to i64
   %mul = mul i64 %sext, -32767
   ret i64 %mul
 }
 
 define <2 x i64> @test9_splat(<2 x i32> %V) {
 ; CHECK-LABEL: @test9_splat(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i32> [[ASHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw <2 x i64> [[SEXT]], <i64 -32767, i64 -32767>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw <2 x i32> [[ASHR]], <i32 -32767, i32 -32767>
+; CHECK-NEXT:    [[MUL:%.*]] = sext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %ashr = ashr <2 x i32> %V, <i32 16, i32 16>
   %sext = sext <2 x i32> %ashr to <2 x i64>
   %mul = mul <2 x i64> %sext, <i64 -32767, i64 -32767>
   ret <2 x i64> %mul
 }
 
 define <2 x i64> @test9_vec(<2 x i32> %V) {
 ; CHECK-LABEL: @test9_vec(
 ; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i32> [[ASHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw <2 x i64> [[SEXT]], <i64 -32767, i64 -10>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nsw <2 x i32> [[ASHR]], <i32 -32767, i32 -10>
+; CHECK-NEXT:    [[MUL:%.*]] = sext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %ashr = ashr <2 x i32> %V, <i32 16, i32 16>
   %sext = sext <2 x i32> %ashr to <2 x i64>
   %mul = mul <2 x i64> %sext, <i64 -32767, i64 -10>
   ret <2 x i64> %mul
 }
 
 define i64 @test10(i32 %V) {
 ; CHECK-LABEL: @test10(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[V:%.*]], 16
 ; CHECK-NEXT:    [[MULCONV:%.*]] = mul nuw i32 [[LSHR]], 65535
 ; CHECK-NEXT:    [[MUL:%.*]] = zext i32 [[MULCONV]] to i64
 ; CHECK-NEXT:    ret i64 [[MUL]]
 ;
   %lshr = lshr i32 %V, 16
   %zext = zext i32 %lshr to i64
   %mul = mul i64 %zext, 65535
   ret i64 %mul
 }
 
 define <2 x i64> @test10_splat(<2 x i32> %V) {
 ; CHECK-LABEL: @test10_splat(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i32> [[LSHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw <2 x i64> [[ZEXT]], <i64 65535, i64 65535>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nuw <2 x i32> [[LSHR]], <i32 65535, i32 65535>
+; CHECK-NEXT:    [[MUL:%.*]] = zext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %lshr = lshr <2 x i32> %V, <i32 16, i32 16>
   %zext = zext <2 x i32> %lshr to <2 x i64>
   %mul = mul <2 x i64> %zext, <i64 65535, i64 65535>
   ret <2 x i64> %mul
 }
 
 define <2 x i64> @test10_vec(<2 x i32> %V) {
 ; CHECK-LABEL: @test10_vec(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[V:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i32> [[LSHR]] to <2 x i64>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw <2 x i64> [[ZEXT]], <i64 65535, i64 2>
+; CHECK-NEXT:    [[MULCONV:%.*]] = mul nuw <2 x i32> [[LSHR]], <i32 65535, i32 2>
+; CHECK-NEXT:    [[MUL:%.*]] = zext <2 x i32> [[MULCONV]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[MUL]]
 ;
   %lshr = lshr <2 x i32> %V, <i32 16, i32 16>
   %zext = zext <2 x i32> %lshr to <2 x i64>
   %mul = mul <2 x i64> %zext, <i64 65535, i64 2>
   ret <2 x i64> %mul
 }
 
 !0 = !{ i32 0, i32 2000 }