[InstCombine] Improve transforms for `(mul X, Y)` -> `(shl X, log2(Y)`

[goldstein.w.n]

Summary

Using the more robust log2 search allows us to fold more cases (same
logic as exists for idiv/irem).

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

489

We should preserve the nuw flag here: https://alive2.llvm.org/ce/z/YkwtoD

We can't preserve nsw due to the INT_MIN edge case: https://alive2.llvm.org/ce/z/HTaC3b

[goldstein.w.n]

Rebase

[goldstein.w.n]

Propegate nuw flag

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

489

We should preserve the nuw flag here: https://alive2.llvm.org/ce/z/YkwtoD

We can't preserve nsw due to the INT_MIN edge case: https://alive2.llvm.org/ce/z/HTaC3b

Also added test this.

[nikic]

LGTM

[goldstein.w.n]

Rebase

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (match(Y, m_OneUse(m_Not(m_OneUse(m_Shl(m_AllOnes(), m_Value(Z))))))) {
     Value *FrX = Builder.CreateFreeze(X, X->getName() + ".fr");
     Value *Shl = Builder.CreateShl(FrX, Z, "mulshl");
     return Builder.CreateSub(Shl, FrX, Mul.getName());
   }
 
   return nullptr;
 }
 
+static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
+                       bool DoFold);
+
 Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   if (Value *V =
           simplifyMulInst(Op0, Op1, I.hasNoSignedWrap(), I.hasNoUnsignedWrap(),
                           SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (SimplifyAssociativeOrCommutative(I))
Value *Abs = Builder.CreateBinaryIntrinsic(
         Intrinsic::abs, X, ConstantInt::getBool(I.getContext(), HasNSW));
     Abs->takeName(&I);
     return replaceInstUsesWith(I, Abs);
   }
 
   if (Instruction *Ext = narrowMathIfNoOverflow(I))
     return Ext;
 
+  // (mul Op0 Op1):
+  //    if Log2(Op0) folds away ->
+  //        (shl Op1, Log2(Op0))
+  //    if Log2(Op1) folds away ->
+  //        (shl Op0, Log2(Op1))
+  if (takeLog2(Builder, Op0, /*Depth*/ 0, /*DoFold*/ false)) {
+    Value *Res = takeLog2(Builder, Op0, /*Depth*/ 0, /*DoFold*/ true);
+    BinaryOperator *Shl = BinaryOperator::CreateShl(Op1, Res);
+    // We can only propegate nuw flag.
+    Shl->setHasNoUnsignedWrap(HasNUW);
+    return Shl;
+  }
+  if (takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ false)) {
+    Value *Res = takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ true);
+    BinaryOperator *Shl = BinaryOperator::CreateShl(Op0, Res);
+    // We can only propegate nuw flag.
+    Shl->setHasNoUnsignedWrap(HasNUW);
+    return Shl;
+  }
+
   bool Changed = false;
   if (!HasNSW && willNotOverflowSignedMul(Op0, Op1, I)) {
     Changed = true;
     I.setHasNoSignedWrap(true);
   }
 
   if (!HasNUW && willNotOverflowUnsignedMul(Op0, Op1, I)) {
     Changed = true;

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=instcombine -S < %s | FileCheck %s
 declare void @use_i8(i8)
 define i8 @mul_selectp2_x(i8 %x, i1 %c) {
 ; CHECK-LABEL: @mul_selectp2_x(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 4
-; CHECK-NEXT:    [[R:%.*]] = mul i8 [[S]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], i8 1, i8 2
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = select i1 %c, i8 2, i8 4
   %r = mul i8 %s, %x
   ret i8 %r
 }
 
 
 define i8 @mul_selectp2_x_propegate_nuw(i8 %x, i1 %c) {
 ; CHECK-LABEL: @mul_selectp2_x_propegate_nuw(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 4
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i8 [[S]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], i8 1, i8 2
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = select i1 %c, i8 2, i8 4
   %r = mul nuw nsw i8 %s, %x
   ret i8 %r
 }
 
 
 define i8 @mul_selectp2_x_multiuse_fixme(i8 %x, i1 %c) {
 ; CHECK-LABEL: @mul_selectp2_x_multiuse_fixme(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 4
-; CHECK-NEXT:    [[R:%.*]] = mul i8 [[S]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C]], i8 1, i8 2
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    call void @use_i8(i8 [[S]])
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = select i1 %c, i8 2, i8 4
   %r = mul i8 %s, %x
   call void @use_i8(i8 %s)
   ret i8 %r
 }
 
 define i8 @mul_selectp2_x_non_const(i8 %x, i1 %c, i8 %yy) {
 ; CHECK-LABEL: @mul_selectp2_x_non_const(
-; CHECK-NEXT:    [[Y:%.*]] = shl nuw i8 1, [[YY:%.*]]
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = mul i8 [[S]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], i8 1, i8 [[YY:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %y = shl i8 1, %yy
   %s = select i1 %c, i8 2, i8 %y
   %r = mul i8 %s, %x
   ret i8 %r
 }
 
 define i8 @mul_selectp2_x_non_const_multiuse(i8 %x, i1 %c, i8 %yy) {
 ; CHECK-LABEL: @mul_selectp2_x_non_const_multiuse(
 ; CHECK-NEXT:    [[Y:%.*]] = shl nuw i8 1, [[YY:%.*]]
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = mul i8 [[S]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], i8 1, i8 [[YY]]
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    call void @use_i8(i8 [[Y]])
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %y = shl i8 1, %yy
   %s = select i1 %c, i8 2, i8 %y
   %r = mul i8 %s, %x
   call void @use_i8(i8 %y)
   ret i8 %r
 }
 
 define i8 @mul_x_selectp2(i8 %xx, i1 %c) {
 ; CHECK-LABEL: @mul_x_selectp2(
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], [[XX]]
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 8, i8 1
-; CHECK-NEXT:    [[R:%.*]] = mul i8 [[X]], [[S]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], i8 3, i8 0
+; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x = mul i8 %xx, %xx
   %s = select i1 %c, i8 8, i8 1
   %r = mul i8 %x, %s
   ret i8 %r
 }
 
 define i8 @mul_select_nonp2_x_fail(i8 %x, i1 %c) {
 ; CHECK-LABEL: @mul_select_nonp2_x_fail(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], i8 2, i8 5
 ; CHECK-NEXT:    [[R:%.*]] = mul i8 [[S]], [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = select i1 %c, i8 2, i8 5
   %r = mul i8 %s, %x
   ret i8 %r
 }
 
 define <2 x i8> @mul_x_selectp2_vec(<2 x i8> %xx, i1 %c) {
 ; CHECK-LABEL: @mul_x_selectp2_vec(
 ; CHECK-NEXT:    [[X:%.*]] = mul <2 x i8> [[XX:%.*]], [[XX]]
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], <2 x i8> <i8 8, i8 16>, <2 x i8> <i8 4, i8 1>
-; CHECK-NEXT:    [[R:%.*]] = mul <2 x i8> [[X]], [[S]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C:%.*]], <2 x i8> <i8 3, i8 4>, <2 x i8> <i8 2, i8 0>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[X]], [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %x = mul <2 x i8> %xx, %xx
   %s = select i1 %c, <2 x i8> <i8 8, i8 16>, <2 x i8> <i8 4, i8 1>
   %r = mul <2 x i8> %x, %s
   ret <2 x i8> %r
 }