[InstCombine] convert mul by negative-pow2 to negate and shift

[spatel]

Summary

This is an unusual canonicalization because we create an extra instruction, but it's very likely better for analysis and codegen (similar reasoning as D133399).

InstCombine::Negator aggressively creates this kind of multiply from subtract and shift, so this is partly reversing a transform in the motivating case from issue #57576.

We may want to limit Negator more to prevent creating multiplies like this. We've seen regressions from that, but no wins AFAIK.

I don't know how to create a fully general proof for this kind of transform in Alive2, but here's an example with bitwidths similar to one of the regression tests:
https://alive2.llvm.org/ce/z/J3jTjR

[spatel]

Patch updated:
No code changes, but rebased after:
73919a87e9a6f8cfa6
...so the motivating test is reduced to remove the multiply and a wide add. We need some other fold to recognize that the wide sub can be narrowed.

[spatel]

Ping.

[RKSimon]

No objections, but I'm a little worried that the negator will put up a fight and convert this back again.....

[spatel]

In D133667#3803056, @RKSimon wrote:

No objections, but I'm a little worried that the negator will put up a fight and convert this back again.....

Yes - I want to believe that the zext protects us from getting into an infinite loop (we shouldn't widen the sub unless it allows removing some other instructions), but it's hard to guarantee that there isn't some loophole.

[nikic]

Can we revert the problematic negator change first?

[spatel]

In D133667#3803288, @nikic wrote:

Can we revert the problematic negator change first?

Yes. There doesn't appear to be any fallout based on regression tests, but I'd do that in 2 pieces to reduce risk. Step 1 would be:
D134310

[spatel]

We inverted the general negator transform with:
ee0bf6472291
...and I haven't seen any complaints after ~10 days. That also partly resolved the diff for issue #57576. This patch still reduces the IR in that example without changing codegen.

Negator will still create a mul from a shl+negate, but it's less likely that we'll conflict with the proposal here, so I think it's safe to give this a try now.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
 
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
   if (Value *V = SimplifyUsingDistributiveLaws(I))
     return replaceInstUsesWith(I, V);
 
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
-  unsigned BitWidth = I.getType()->getScalarSizeInBits();
+  Type *Ty = I.getType();
+  unsigned BitWidth = Ty->getScalarSizeInBits();
 
   // X * -1 == 0 - X
   if (match(Op1, m_AllOnes())) {
     BinaryOperator *BO = BinaryOperator::CreateNeg(Op0, I.getName());
     if (I.hasNoSignedWrap())
       BO->setHasNoSignedWrap();
     return BO;
   }
Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   }
 
   if (Op0->hasOneUse() && match(Op1, m_NegatedPower2())) {
     // Interpret  X * (-1<<C)  as  (-X) * (1<<C)  and try to sink the negation.
     // The "* (1<<C)" thus becomes a potential shifting opportunity.
     if (Value *NegOp0 = Negator::Negate(/*IsNegation*/ true, Op0, *this))
       return BinaryOperator::CreateMul(
           NegOp0, ConstantExpr::getNeg(cast<Constant>(Op1)), I.getName());
+
+    // Try to convert multiply of extended operand to narrow negate and shift
+    // for better analysis.
+    // This is valid if the shift amount (trailing zeros in the multiplier
+    // constant) clears more high bits than the bitwidth difference between
+    // source and destination types:
+    // ({z/s}ext X) * (-1<<C) --> (zext (-X)) << C
+    const APInt *NegPow2C;
+    Value *X;
+    if (match(Op0, m_ZExtOrSExt(m_Value(X))) &&
+        match(Op1, m_APIntAllowUndef(NegPow2C))) {
+      unsigned SrcWidth = X->getType()->getScalarSizeInBits();
+      unsigned ShiftAmt = NegPow2C->countTrailingZeros();
+      if (ShiftAmt >= BitWidth - SrcWidth) {
+        Value *N = Builder.CreateNeg(X, X->getName() + ".neg");
+        Value *Z = Builder.CreateZExt(N, Ty, N->getName() + ".z");
+        return BinaryOperator::CreateShl(Z, ConstantInt::get(Ty, ShiftAmt));
+      }
+    }
   }
 
   if (Instruction *FoldedMul = foldBinOpIntoSelectOrPhi(I))
     return FoldedMul;
 
   if (Value *FoldedMul = foldMulSelectToNegate(I, Builder))
     return replaceInstUsesWith(I, FoldedMul);
 
// (X / Y) * -Y = (X % Y) - X
     }
   }
 
   // Fold the following two scenarios:
   //   1) i1 mul -> i1 and.
   //   2) X * Y --> X & Y, iff X, Y can be only {0,1}.
   // Note: We could use known bits to generalize this and related patterns with
   // shifts/truncs
-  Type *Ty = I.getType();
   if (Ty->isIntOrIntVectorTy(1) ||
       (match(Op0, m_And(m_Value(), m_One())) &&
        match(Op1, m_And(m_Value(), m_One()))))
     return BinaryOperator::CreateAnd(Op0, Op1);
 
   // X*(1 << Y) --> X << Y
   // (1 << Y)*X --> X << Y
   {

llvm/test/Transforms/InstCombine/add.ll

;
   %r = shl i8 %s, 1
   ret i8 %r
 }
 
 define { i64, i64 } @PR57576(i64 noundef %x, i64 noundef %y, i64 noundef %z, i64 noundef %w) {
 ; CHECK-LABEL: @PR57576(
 ; CHECK-NEXT:    [[ZX:%.*]] = zext i64 [[X:%.*]] to i128
 ; CHECK-NEXT:    [[ZY:%.*]] = zext i64 [[Y:%.*]] to i128
-; CHECK-NEXT:    [[ZW:%.*]] = zext i64 [[W:%.*]] to i128
 ; CHECK-NEXT:    [[ZZ:%.*]] = zext i64 [[Z:%.*]] to i128
 ; CHECK-NEXT:    [[SHY:%.*]] = shl nuw i128 [[ZY]], 64
 ; CHECK-NEXT:    [[XY:%.*]] = or i128 [[SHY]], [[ZX]]
-; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i128 [[ZW]], 64
-; CHECK-NEXT:    [[TMP2:%.*]] = or i128 [[TMP1]], [[ZZ]]
-; CHECK-NEXT:    [[ADD:%.*]] = sub i128 [[XY]], [[TMP2]]
-; CHECK-NEXT:    [[T:%.*]] = trunc i128 [[ADD]] to i64
-; CHECK-NEXT:    [[H:%.*]] = lshr i128 [[ADD]], 64
-; CHECK-NEXT:    [[T2:%.*]] = trunc i128 [[H]] to i64
+; CHECK-NEXT:    [[SUB:%.*]] = sub i128 [[XY]], [[ZZ]]
+; CHECK-NEXT:    [[T:%.*]] = trunc i128 [[SUB]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i128 [[SUB]], 64
+; CHECK-NEXT:    [[DOTTR:%.*]] = trunc i128 [[TMP1]] to i64
+; CHECK-NEXT:    [[DOTNARROW:%.*]] = sub i64 [[DOTTR]], [[W:%.*]]
 ; CHECK-NEXT:    [[R1:%.*]] = insertvalue { i64, i64 } poison, i64 [[T]], 0
-; CHECK-NEXT:    [[R2:%.*]] = insertvalue { i64, i64 } [[R1]], i64 [[T2]], 1
+; CHECK-NEXT:    [[R2:%.*]] = insertvalue { i64, i64 } [[R1]], i64 [[DOTNARROW]], 1
 ; CHECK-NEXT:    ret { i64, i64 } [[R2]]
 ;
   %zx = zext i64 %x to i128
   %zy = zext i64 %y to i128
   %zw = zext i64 %w to i128
   %zz = zext i64 %z to i128
   %shy = shl nuw i128 %zy, 64
   %mw = mul i128 %zw, -18446744073709551616

llvm/test/Transforms/InstCombine/mul.ll

;
   %not = xor i32 %a0, -1
   call void @use32(i32 %not)
   %mul = mul i32 %not, -4
   ret i32 %mul
 }
 
 define i32 @zext_negpow2(i8 %x) {
 ; CHECK-LABEL: @zext_negpow2(
-; CHECK-NEXT:    [[ZX:%.*]] = zext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[ZX]], -16777216
+; CHECK-NEXT:    [[X_NEG:%.*]] = sub i8 0, [[X:%.*]]
+; CHECK-NEXT:    [[X_NEG_Z:%.*]] = zext i8 [[X_NEG]] to i32
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i32 [[X_NEG_Z]], 24
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %zx = zext i8 %x to i32
   %r = mul i32 %zx, -16777216 ; -1 << 24
   ret i32 %r
 }
 
+; splat constant
+
 define <2 x i14> @zext_negpow2_vec(<2 x i5> %x) {
 ; CHECK-LABEL: @zext_negpow2_vec(
-; CHECK-NEXT:    [[ZX:%.*]] = zext <2 x i5> [[X:%.*]] to <2 x i14>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i14> [[ZX]], <i14 -2048, i14 -2048>
+; CHECK-NEXT:    [[X_NEG:%.*]] = sub <2 x i5> zeroinitializer, [[X:%.*]]
+; CHECK-NEXT:    [[X_NEG_Z:%.*]] = zext <2 x i5> [[X_NEG]] to <2 x i14>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i14> [[X_NEG_Z]], <i14 11, i14 11>
 ; CHECK-NEXT:    ret <2 x i14> [[R]]
 ;
   %zx = zext <2 x i5> %x to <2 x i14>
   %r = mul <2 x i14> %zx, <i14 -2048, i14 -2048> ; -1 << 11
   ret <2 x i14> %r
 }
 
+; negative test - mul must be big enough to cover bitwidth diff
+
 define i32 @zext_negpow2_too_small(i8 %x) {
 ; CHECK-LABEL: @zext_negpow2_too_small(
 ; CHECK-NEXT:    [[ZX:%.*]] = zext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[R:%.*]] = mul nsw i32 [[ZX]], -8388608
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %zx = zext i8 %x to i32
   %r = mul i32 %zx, -8388608 ; -1 << 23
   ret i32 %r
 }
 
 define i16 @sext_negpow2(i9 %x) {
 ; CHECK-LABEL: @sext_negpow2(
-; CHECK-NEXT:    [[SX:%.*]] = sext i9 [[X:%.*]] to i16
-; CHECK-NEXT:    [[R:%.*]] = mul i16 [[SX]], -1024
+; CHECK-NEXT:    [[X_NEG:%.*]] = sub i9 0, [[X:%.*]]
+; CHECK-NEXT:    [[X_NEG_Z:%.*]] = zext i9 [[X_NEG]] to i16
+; CHECK-NEXT:    [[R:%.*]] = shl i16 [[X_NEG_Z]], 10
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %sx = sext i9 %x to i16
   %r = mul i16 %sx, -1024 ; -1 << 10
   ret i16 %r
 }
 
+; splat constant with poison element(s)
+
 define <2 x i16> @sext_negpow2_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @sext_negpow2_vec(
-; CHECK-NEXT:    [[SX:%.*]] = sext <2 x i8> [[X:%.*]] to <2 x i16>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i16> [[SX]], <i16 -256, i16 poison>
+; CHECK-NEXT:    [[X_NEG:%.*]] = sub <2 x i8> zeroinitializer, [[X:%.*]]
+; CHECK-NEXT:    [[X_NEG_Z:%.*]] = zext <2 x i8> [[X_NEG]] to <2 x i16>
+; CHECK-NEXT:    [[R:%.*]] = shl nuw <2 x i16> [[X_NEG_Z]], <i16 8, i16 8>
 ; CHECK-NEXT:    ret <2 x i16> [[R]]
 ;
   %sx = sext <2 x i8> %x to <2 x i16>
   %r = mul <2 x i16> %sx, <i16 -256, i16 poison> ; -1 << 8
   ret <2 x i16> %r
 }
 
+; negative test - mul must be big enough to cover bitwidth diff
+
 define <2 x i16> @sext_negpow2_too_small_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @sext_negpow2_too_small_vec(
 ; CHECK-NEXT:    [[SX:%.*]] = sext <2 x i8> [[X:%.*]] to <2 x i16>
 ; CHECK-NEXT:    [[R:%.*]] = mul <2 x i16> [[SX]], <i16 -128, i16 poison>
 ; CHECK-NEXT:    ret <2 x i16> [[R]]
 ;
   %sx = sext <2 x i8> %x to <2 x i16>
   %r = mul <2 x i16> %sx, <i16 -128, i16 poison> ; -1 << 7
   ret <2 x i16> %r
 }
 
+; negative test - too many uses
+
 define i32 @zext_negpow2_use(i8 %x) {
 ; CHECK-LABEL: @zext_negpow2_use(
 ; CHECK-NEXT:    [[ZX:%.*]] = zext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    call void @use32(i32 [[ZX]])
 ; CHECK-NEXT:    [[R:%.*]] = mul i32 [[ZX]], -16777216
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %zx = zext i8 %x to i32
   call void @use32(i32 %zx)
   %r = mul i32 %zx, -16777216 ; -1 << 24
   ret i32 %r
 }