[InstCombine] Fold mul (lshr exact (X, 2^N + 1)), N -> add (X , lshr (X, N))

Alive2 Proofs:
https://alive2.llvm.org/ce/z/LVqGEo
https://alive2.llvm.org/ce/z/dyeGEv

Author: AZero13

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

@@ -257,6 +257,35 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
     }
   }
 
+  {
+    // mul (lshr exact X, N), (2^N + 1) -> add (X, lshr exact (X, N))
+    Value *NewOp;
+    const APInt *ShiftC;
+    const APInt *MulAP;
+    if (match(&I, m_Mul(m_Exact(m_Shr(m_Value(NewOp), m_APInt(ShiftC))),
+                        m_APInt(MulAP)))) {
+      if (BitWidth > 2 && (*MulAP - 1).isPowerOf2() &&
+          *ShiftC == MulAP->logBase2()) {
+        Value *BinOp = Op0;
+        BinaryOperator *OpBO = cast<BinaryOperator>(Op0);
+        if (!HasNUW && !HasNSW)
+          NewOp = Builder.CreateFreeze(NewOp, ".fr");
+        else if (HasNUW && OpBO->getOpcode() == Instruction::AShr &&
+                 OpBO->hasOneUse())
+          BinOp = Builder.CreateLShr(NewOp, ConstantInt::get(Ty, *ShiftC), "",
+                                     /*isExact=*/true);
+
+        auto *NewAdd = BinaryOperator::CreateAdd(NewOp, BinOp);
+        if (HasNSW && (OpBO->getOpcode() == Instruction::LShr ||
+                       ShiftC->getZExtValue() < BitWidth - 1))
+          NewAdd->setHasNoSignedWrap(true);
+
+        NewAdd->setHasNoUnsignedWrap(HasNUW);
+        return NewAdd;
+      }
+    }
+  }
+
   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.

llvm/test/Transforms/InstCombine/ashr-lshr.ll

@@ -864,8 +864,9 @@ define i32 @ashr_mul_times_5_div_4_exact_2(i32 %x) {
 
 define i32 @ashr_shift_mul(i32 %x) {
 ; CHECK-LABEL: @ashr_shift_mul(
-; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul i32 [[A]], 9
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[DOTFR]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[DOTFR]], [[A]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = ashr exact i32 %x, 3
@@ -875,8 +876,8 @@ define i32 @ashr_shift_mul(i32 %x) {
 
 define i32 @ashr_shift_mul_nuw(i32 %x) {
 ; CHECK-LABEL: @ashr_shift_mul_nuw(
-; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul nuw i32 [[A]], 9
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = ashr exact i32 %x, 3
@@ -887,7 +888,7 @@ define i32 @ashr_shift_mul_nuw(i32 %x) {
 define i32 @ashr_shift_mul_nsw(i32 %x) {
 ; CHECK-LABEL: @ashr_shift_mul_nsw(
 ; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul nsw i32 [[A]], 9
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[A]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = ashr exact i32 %x, 3
@@ -898,7 +899,7 @@ define i32 @ashr_shift_mul_nsw(i32 %x) {
 define i32 @lshr_shift_mul_nuw(i32 %x) {
 ; CHECK-LABEL: @lshr_shift_mul_nuw(
 ; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul nuw i32 [[A]], 9
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[A]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = lshr exact i32 %x, 3
@@ -908,8 +909,9 @@ define i32 @lshr_shift_mul_nuw(i32 %x) {
 
 define i32 @lshr_shift_mul(i32 %x) {
 ; CHECK-LABEL: @lshr_shift_mul(
-; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul i32 [[A]], 9
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[DOTFR]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[DOTFR]], [[A]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = lshr exact i32 %x, 3
@@ -920,7 +922,7 @@ define i32 @lshr_shift_mul(i32 %x) {
 define i32 @lshr_shift_mul_nsw(i32 %x) {
 ; CHECK-LABEL: @lshr_shift_mul_nsw(
 ; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[RES:%.*]] = mul nuw nsw i32 [[A]], 9
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[A]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = lshr exact i32 %x, 3
@@ -954,13 +956,11 @@ define i32 @ashr_no_exact(i32 %x) {
   ret i32 %res
 }
 
-; Negative test
-
 define i32 @lshr_multiuse(i32 %x) {
 ; CHECK-LABEL: @lshr_multiuse(
 ; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
 ; CHECK-NEXT:    call void @use(i32 [[A]])
-; CHECK-NEXT:    [[RES:%.*]] = mul nuw nsw i32 [[A]], 9
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[A]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = lshr exact i32 %x, 3
@@ -997,11 +997,41 @@ define i32 @ashr_multiuse_no_flags(i32 %x) {
   ret i32 %res
 }
 
+define i32 @lshr_multiuse_no_flags(i32 %x) {
+; CHECK-LABEL: @lshr_multiuse_no_flags(
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[DOTFR]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[DOTFR]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  call void @use(i32 %a)
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
+; Negative test
+
+define i32 @ashr_multiuse_no_flags(i32 %x) {
+; CHECK-LABEL: @ashr_multiuse_no_flags(
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[DOTFR]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[DOTFR]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  call void @use(i32 %a)
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
 define i32 @ashr_multiuse(i32 %x) {
 ; CHECK-LABEL: @ashr_multiuse(
 ; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
 ; CHECK-NEXT:    call void @use(i32 [[A]])
-; CHECK-NEXT:    [[RES:%.*]] = mul nsw i32 [[A]], 9
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[A]], [[X]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %a = ashr exact i32 %x, 3