Skip to content

Commit

Permalink
[InstCombine] Fold mul (shr exact (X, N)), 2^N + 1 -> add (X , shr ex…
Browse files Browse the repository at this point in the history 
…act (X, N))

Alive2 Proofs:
https://alive2.llvm.org/ce/z/aJnxyp
https://alive2.llvm.org/ce/z/dyeGEv
  • Loading branch information
@AZero13
AZero13 committed Nov 14, 2024
1 parent 5bde6a4 commit 37b0139
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 15 deletions.
31 changes: 31 additions & 0 deletions llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,37 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
}
}

// mul (shr exact X, N), (2^N + 1) -> add (X, shr 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");

// mul (ashr nuw exact X, N) -> add (X, lshr nuw exact (X, N))
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.
Expand Down
34 changes: 19 additions & 15 deletions llvm/test/Transforms/InstCombine/ashr-lshr.ll
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1075,10 +1075,12 @@ entry:
call void @use(i32 %and)
%shr = ashr i32 %and, 31
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:%.*]] = mul i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = ashr exact i32 %x, 3
Expand All @@ -1088,8 +1090,8 @@ define i32 @ashr_shift_mul(i32 noundef %x) {

define i32 @ashr_shift_mul_nuw(i32 noundef %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 [[X]], [[TMP1]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = ashr exact i32 %x, 3
Expand All @@ -1100,7 +1102,7 @@ define i32 @ashr_shift_mul_nuw(i32 noundef %x) {
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:%.*]] = mul nsw i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add nsw i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = ashr exact i32 %x, 3
Expand All @@ -1111,7 +1113,7 @@ define i32 @ashr_shift_mul_nsw(i32 noundef %x) {
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:%.*]] = mul nuw i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add nuw i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = lshr exact i32 %x, 3
Expand All @@ -1122,7 +1124,7 @@ define i32 @lshr_shift_mul_nuw(i32 noundef %x) {
define i32 @lshr_shift_mul(i32 noundef %x) {
; CHECK-LABEL: @lshr_shift_mul(
; CHECK-NEXT: [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
; CHECK-NEXT: [[RES:%.*]] = mul i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = lshr exact i32 %x, 3
Expand All @@ -1133,7 +1135,7 @@ define i32 @lshr_shift_mul(i32 noundef %x) {
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:%.*]] = mul nuw nsw i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add nsw i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = lshr exact i32 %x, 3
Expand Down Expand Up @@ -1169,8 +1171,9 @@ define i32 @ashr_no_exact(i32 %x) {

define i32 @lshr_no_undef(i32 %x) {
; CHECK-LABEL: @lshr_no_undef(
; CHECK-NEXT: [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
; CHECK-NEXT: [[RES:%.*]] = mul nuw nsw i32 [[A]], 9
; 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
Expand All @@ -1180,8 +1183,9 @@ define i32 @lshr_no_undef(i32 %x) {

define i32 @ashr_no_undef(i32 %x) {
; CHECK-LABEL: @ashr_no_undef(
; CHECK-NEXT: [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
; CHECK-NEXT: [[RES:%.*]] = mul nsw i32 [[A]], 9
; 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
Expand All @@ -1193,7 +1197,7 @@ 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:%.*]] = mul nuw nsw i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add nsw i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = lshr exact i32 %x, 3
Expand All @@ -1206,7 +1210,7 @@ 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:%.*]] = mul i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = lshr exact i32 %x, 3
Expand All @@ -1219,7 +1223,7 @@ 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:%.*]] = mul i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = ashr exact i32 %x, 3
Expand All @@ -1232,7 +1236,7 @@ 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:%.*]] = mul nsw i32 [[A]], 9
; CHECK-NEXT: [[RES:%.*]] = add nsw i32 [[X]], [[A]]
; CHECK-NEXT: ret i32 [[RES]]
;
%a = ashr exact i32 %x, 3
Expand Down

0 comments on commit 37b0139

Please sign in to comment.