Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

AZero13
Copy link
Contributor

@AZero13 AZero13 commented Oct 15, 2024

@llvmbot
Copy link
Member

llvmbot commented Oct 15, 2024

@llvm/pr-subscribers-llvm-transforms

Author: Rose (AreaZR)

Changes

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


Full diff: https://github.com/llvm/llvm-project/pull/112407.diff

2 Files Affected:

  • (modified) llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp (+29)
  • (modified) llvm/test/Transforms/InstCombine/ashr-lshr.ll (+168)
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index f4f3644acfe5ea..3378da1130ae68 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -261,6 +261,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 (!isGuaranteedNotToBeUndef(NewOp, &AC, &I, &DT))
+          NewOp = Builder.CreateFreeze(NewOp, NewOp->getName() + ".fr");
+        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 && (HasNUW || 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.
diff --git a/llvm/test/Transforms/InstCombine/ashr-lshr.ll b/llvm/test/Transforms/InstCombine/ashr-lshr.ll
index 1abf1be2cbedd9..9c4611d0be3947 100644
--- a/llvm/test/Transforms/InstCombine/ashr-lshr.ll
+++ b/llvm/test/Transforms/InstCombine/ashr-lshr.ll
@@ -1077,4 +1077,172 @@ entry:
   ret i32 %shr
 }
 
+define i32 @ashr_shift_mul(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @ashr_shift_mul_nuw(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul_nuw(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[X]], [[TMP1]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nuw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @ashr_shift_mul_nsw(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul_nsw(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul_nuw(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul_nuw(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul nuw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul_nsw(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul_nsw(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+; Negative test
+
+define i32 @lshr_no_exact(i32 %x) {
+; CHECK-LABEL: @lshr_no_exact(
+; CHECK-NEXT:    [[A:%.*]] = lshr i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = mul nuw nsw i32 [[A]], 9
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+; Negative test
+
+define i32 @ashr_no_exact(i32 %x) {
+; CHECK-LABEL: @ashr_no_exact(
+; CHECK-NEXT:    [[A:%.*]] = ashr i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = mul nsw i32 [[A]], 9
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_no_undef(i32 %x) {
+; CHECK-LABEL: @lshr_no_undef(
+; CHECK-NEXT:    [[X_FR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X_FR]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X_FR]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @ashr_no_undef(i32 %x) {
+; CHECK-LABEL: @ashr_no_undef(
+; CHECK-NEXT:    [[X_FR:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X_FR]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X_FR]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_multiuse(i32 noundef %x) {
+; CHECK-LABEL: @lshr_multiuse(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  call void @use(i32 %a)
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_multiuse_no_flags(i32 noundef %x) {
+; CHECK-LABEL: @lshr_multiuse_no_flags(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[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
+}
+
+define i32 @ashr_multiuse_no_flags(i32 noundef %x) {
+; CHECK-LABEL: @ashr_multiuse_no_flags(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[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 noundef %x) {
+; CHECK-LABEL: @ashr_multiuse(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    call void @use(i32 [[A]])
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  call void @use(i32 %a)
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
 declare void @use(i32)

@AZero13 AZero13 changed the title [InstCombine] Fold mul (lshr exact (X, N)), 2^N + 1 -> add (X , lshr exact (X, N)) [InstCombine] Fold mul (shr exact (X, N)), 2^N + 1 -> add (X , shr exact (X, N)) Oct 15, 2024
@AZero13 AZero13 force-pushed the inverse-div-shift branch 2 times, most recently from c6a0f7e to 35698c1 Compare October 15, 2024 17:46
@AZero13
Copy link
Contributor Author

AZero13 commented Nov 17, 2024

Going to ping @nikic

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants