Skip to content

Commit

Permalink
Reapply [InstCombine] lshr (mul (X, 2^N + 1)), N -> add (X, lshr(X, N…
Browse files Browse the repository at this point in the history 
…)) (#92907)

Alive2 Proofs:
https://alive2.llvm.org/ce/z/eSinJY
https://alive2.llvm.org/ce/z/vyKvde
https://alive2.llvm.org/ce/z/dRFsfV

I mistakenly reverted this commit as part of a larger set of
reverts. Reapplied without changes.
  • Loading branch information
@AreaZR @nikic
AreaZR authored and nikic committed May 29, 2024
1 parent 3bcccb6 commit 718ba5a
Show file tree
Hide file tree
Showing 3 changed files with 318 additions and 10 deletions.
50 changes: 43 additions & 7 deletions llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1461,13 +1461,24 @@ Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {

const APInt *MulC;
if (match(Op0, m_NUWMul(m_Value(X), m_APInt(MulC)))) {
// Look for a "splat" mul pattern - it replicates bits across each half of
// a value, so a right shift is just a mask of the low bits:
// lshr i[2N] (mul nuw X, (2^N)+1), N --> and iN X, (2^N)-1
// TODO: Generalize to allow more than just half-width shifts?
if (BitWidth > 2 && ShAmtC * 2 == BitWidth && (*MulC - 1).isPowerOf2() &&
MulC->logBase2() == ShAmtC)
return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, *MulC - 2));
if (BitWidth > 2 && (*MulC - 1).isPowerOf2() &&
MulC->logBase2() == ShAmtC) {
// Look for a "splat" mul pattern - it replicates bits across each half
// of a value, so a right shift is just a mask of the low bits:
// lshr i[2N] (mul nuw X, (2^N)+1), N --> and iN X, (2^N)-1
if (ShAmtC * 2 == BitWidth)
return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, *MulC - 2));

// lshr (mul nuw (X, 2^N + 1)), N -> add nuw (X, lshr(X, N))
if (Op0->hasOneUse()) {
auto *NewAdd = BinaryOperator::CreateNUWAdd(
X, Builder.CreateLShr(X, ConstantInt::get(Ty, ShAmtC), "",
I.isExact()));
NewAdd->setHasNoSignedWrap(
cast<OverflowingBinaryOperator>(Op0)->hasNoSignedWrap());
return NewAdd;
}
}

// The one-use check is not strictly necessary, but codegen may not be
// able to invert the transform and perf may suffer with an extra mul
Expand All @@ -1487,6 +1498,16 @@ Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
}
}

// lshr (mul nsw (X, 2^N + 1)), N -> add nsw (X, lshr(X, N))
if (match(Op0, m_OneUse(m_NSWMul(m_Value(X), m_APInt(MulC))))) {
if (BitWidth > 2 && (*MulC - 1).isPowerOf2() &&
MulC->logBase2() == ShAmtC) {
return BinaryOperator::CreateNSWAdd(
X, Builder.CreateLShr(X, ConstantInt::get(Ty, ShAmtC), "",
I.isExact()));
}
}

// Try to narrow bswap.
// In the case where the shift amount equals the bitwidth difference, the
// shift is eliminated.
Expand Down Expand Up @@ -1690,6 +1711,21 @@ Instruction *InstCombinerImpl::visitAShr(BinaryOperator &I) {
if (match(Op0, m_OneUse(m_NSWSub(m_Value(X), m_Value(Y)))))
return new SExtInst(Builder.CreateICmpSLT(X, Y), Ty);
}

const APInt *MulC;
if (match(Op0, m_OneUse(m_NSWMul(m_Value(X), m_APInt(MulC)))) &&
(BitWidth > 2 && (*MulC - 1).isPowerOf2() &&
MulC->logBase2() == ShAmt &&
(ShAmt < BitWidth - 1))) /* Minus 1 for the sign bit */ {

// ashr (mul nsw (X, 2^N + 1)), N -> add nsw (X, ashr(X, N))
auto *NewAdd = BinaryOperator::CreateNSWAdd(
X,
Builder.CreateAShr(X, ConstantInt::get(Ty, ShAmt), "", I.isExact()));
NewAdd->setHasNoUnsignedWrap(
cast<OverflowingBinaryOperator>(Op0)->hasNoUnsignedWrap());
return NewAdd;
}
}

const SimplifyQuery Q = SQ.getWithInstruction(&I);
Expand Down
259 changes: 259 additions & 0 deletions llvm/test/Transforms/InstCombine/ashr-lshr.ll
View file
Original file line number Diff line number Diff line change
Expand Up @@ -604,3 +604,262 @@ define <2 x i8> @ashr_known_pos_exact_vec(<2 x i8> %x, <2 x i8> %y) {
%r = ashr exact <2 x i8> %p, %y
ret <2 x i8> %r
}

define i32 @lshr_mul_times_3_div_2(i32 %0) {
; CHECK-LABEL: @lshr_mul_times_3_div_2(
; CHECK-NEXT: [[TMP2:%.*]] = lshr i32 [[TMP0:%.*]], 1
; CHECK-NEXT: [[LSHR:%.*]] = add nuw nsw i32 [[TMP2]], [[TMP0]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nsw nuw i32 %0, 3
%lshr = lshr i32 %mul, 1
ret i32 %lshr
}

define i32 @lshr_mul_times_3_div_2_exact(i32 %x) {
; CHECK-LABEL: @lshr_mul_times_3_div_2_exact(
; CHECK-NEXT: [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 1
; CHECK-NEXT: [[LSHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nsw i32 %x, 3
%lshr = lshr exact i32 %mul, 1
ret i32 %lshr
}

; Negative test

define i32 @lshr_mul_times_3_div_2_no_flags(i32 %0) {
; CHECK-LABEL: @lshr_mul_times_3_div_2_no_flags(
; CHECK-NEXT: [[MUL:%.*]] = mul i32 [[TMP0:%.*]], 3
; CHECK-NEXT: [[LSHR:%.*]] = lshr i32 [[MUL]], 1
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul i32 %0, 3
%lshr = lshr i32 %mul, 1
ret i32 %lshr
}

; Negative test

define i32 @mul_times_3_div_2_multiuse_lshr(i32 %x) {
; CHECK-LABEL: @mul_times_3_div_2_multiuse_lshr(
; CHECK-NEXT: [[MUL:%.*]] = mul nuw i32 [[X:%.*]], 3
; CHECK-NEXT: [[RES:%.*]] = lshr i32 [[MUL]], 1
; CHECK-NEXT: call void @use(i32 [[MUL]])
; CHECK-NEXT: ret i32 [[RES]]
;
%mul = mul nuw i32 %x, 3
%res = lshr i32 %mul, 1
call void @use(i32 %mul)
ret i32 %res
}

define i32 @lshr_mul_times_3_div_2_exact_2(i32 %x) {
; CHECK-LABEL: @lshr_mul_times_3_div_2_exact_2(
; CHECK-NEXT: [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 1
; CHECK-NEXT: [[LSHR:%.*]] = add nuw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nuw i32 %x, 3
%lshr = lshr exact i32 %mul, 1
ret i32 %lshr
}

define i32 @lshr_mul_times_5_div_4(i32 %0) {
; CHECK-LABEL: @lshr_mul_times_5_div_4(
; CHECK-NEXT: [[TMP2:%.*]] = lshr i32 [[TMP0:%.*]], 2
; CHECK-NEXT: [[LSHR:%.*]] = add nuw nsw i32 [[TMP2]], [[TMP0]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nsw nuw i32 %0, 5
%lshr = lshr i32 %mul, 2
ret i32 %lshr
}

define i32 @lshr_mul_times_5_div_4_exact(i32 %x) {
; CHECK-LABEL: @lshr_mul_times_5_div_4_exact(
; CHECK-NEXT: [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 2
; CHECK-NEXT: [[LSHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nsw i32 %x, 5
%lshr = lshr exact i32 %mul, 2
ret i32 %lshr
}

; Negative test

define i32 @lshr_mul_times_5_div_4_no_flags(i32 %0) {
; CHECK-LABEL: @lshr_mul_times_5_div_4_no_flags(
; CHECK-NEXT: [[MUL:%.*]] = mul i32 [[TMP0:%.*]], 5
; CHECK-NEXT: [[LSHR:%.*]] = lshr i32 [[MUL]], 2
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul i32 %0, 5
%lshr = lshr i32 %mul, 2
ret i32 %lshr
}

; Negative test

define i32 @mul_times_5_div_4_multiuse_lshr(i32 %x) {
; CHECK-LABEL: @mul_times_5_div_4_multiuse_lshr(
; CHECK-NEXT: [[MUL:%.*]] = mul nuw i32 [[X:%.*]], 5
; CHECK-NEXT: [[RES:%.*]] = lshr i32 [[MUL]], 2
; CHECK-NEXT: call void @use(i32 [[MUL]])
; CHECK-NEXT: ret i32 [[RES]]
;
%mul = mul nuw i32 %x, 5
%res = lshr i32 %mul, 2
call void @use(i32 %mul)
ret i32 %res
}

define i32 @lshr_mul_times_5_div_4_exact_2(i32 %x) {
; CHECK-LABEL: @lshr_mul_times_5_div_4_exact_2(
; CHECK-NEXT: [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 2
; CHECK-NEXT: [[LSHR:%.*]] = add nuw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[LSHR]]
;
%mul = mul nuw i32 %x, 5
%lshr = lshr exact i32 %mul, 2
ret i32 %lshr
}

define i32 @ashr_mul_times_3_div_2(i32 %0) {
; CHECK-LABEL: @ashr_mul_times_3_div_2(
; CHECK-NEXT: [[TMP2:%.*]] = ashr i32 [[TMP0:%.*]], 1
; CHECK-NEXT: [[ASHR:%.*]] = add nuw nsw i32 [[TMP2]], [[TMP0]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nuw nsw i32 %0, 3
%ashr = ashr i32 %mul, 1
ret i32 %ashr
}

define i32 @ashr_mul_times_3_div_2_exact(i32 %x) {
; CHECK-LABEL: @ashr_mul_times_3_div_2_exact(
; CHECK-NEXT: [[TMP1:%.*]] = ashr exact i32 [[X:%.*]], 1
; CHECK-NEXT: [[ASHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nsw i32 %x, 3
%ashr = ashr exact i32 %mul, 1
ret i32 %ashr
}

; Negative test

define i32 @ashr_mul_times_3_div_2_no_flags(i32 %0) {
; CHECK-LABEL: @ashr_mul_times_3_div_2_no_flags(
; CHECK-NEXT: [[MUL:%.*]] = mul i32 [[TMP0:%.*]], 3
; CHECK-NEXT: [[ASHR:%.*]] = ashr i32 [[MUL]], 1
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul i32 %0, 3
%ashr = ashr i32 %mul, 1
ret i32 %ashr
}

; Negative test

define i32 @ashr_mul_times_3_div_2_no_nsw(i32 %0) {
; CHECK-LABEL: @ashr_mul_times_3_div_2_no_nsw(
; CHECK-NEXT: [[MUL:%.*]] = mul nuw i32 [[TMP0:%.*]], 3
; CHECK-NEXT: [[ASHR:%.*]] = ashr i32 [[MUL]], 1
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nuw i32 %0, 3
%ashr = ashr i32 %mul, 1
ret i32 %ashr
}

; Negative test

define i32 @mul_times_3_div_2_multiuse_ashr(i32 %x) {
; CHECK-LABEL: @mul_times_3_div_2_multiuse_ashr(
; CHECK-NEXT: [[MUL:%.*]] = mul nsw i32 [[X:%.*]], 3
; CHECK-NEXT: [[RES:%.*]] = ashr i32 [[MUL]], 1
; CHECK-NEXT: call void @use(i32 [[MUL]])
; CHECK-NEXT: ret i32 [[RES]]
;
%mul = mul nsw i32 %x, 3
%res = ashr i32 %mul, 1
call void @use(i32 %mul)
ret i32 %res
}

define i32 @ashr_mul_times_3_div_2_exact_2(i32 %x) {
; CHECK-LABEL: @ashr_mul_times_3_div_2_exact_2(
; CHECK-NEXT: [[TMP1:%.*]] = ashr exact i32 [[X:%.*]], 1
; CHECK-NEXT: [[ASHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nsw i32 %x, 3
%ashr = ashr exact i32 %mul, 1
ret i32 %ashr
}

define i32 @ashr_mul_times_5_div_4(i32 %0) {
; CHECK-LABEL: @ashr_mul_times_5_div_4(
; CHECK-NEXT: [[TMP2:%.*]] = ashr i32 [[TMP0:%.*]], 2
; CHECK-NEXT: [[ASHR:%.*]] = add nuw nsw i32 [[TMP2]], [[TMP0]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nuw nsw i32 %0, 5
%ashr = ashr i32 %mul, 2
ret i32 %ashr
}

define i32 @ashr_mul_times_5_div_4_exact(i32 %x) {
; CHECK-LABEL: @ashr_mul_times_5_div_4_exact(
; CHECK-NEXT: [[TMP1:%.*]] = ashr exact i32 [[X:%.*]], 2
; CHECK-NEXT: [[ASHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nsw i32 %x, 5
%ashr = ashr exact i32 %mul, 2
ret i32 %ashr
}

; Negative test

define i32 @ashr_mul_times_5_div_4_no_flags(i32 %0) {
; CHECK-LABEL: @ashr_mul_times_5_div_4_no_flags(
; CHECK-NEXT: [[MUL:%.*]] = mul i32 [[TMP0:%.*]], 5
; CHECK-NEXT: [[ASHR:%.*]] = ashr i32 [[MUL]], 2
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul i32 %0, 5
%ashr = ashr i32 %mul, 2
ret i32 %ashr
}

; Negative test

define i32 @mul_times_5_div_4_multiuse_ashr(i32 %x) {
; CHECK-LABEL: @mul_times_5_div_4_multiuse_ashr(
; CHECK-NEXT: [[MUL:%.*]] = mul nsw i32 [[X:%.*]], 5
; CHECK-NEXT: [[RES:%.*]] = ashr i32 [[MUL]], 2
; CHECK-NEXT: call void @use(i32 [[MUL]])
; CHECK-NEXT: ret i32 [[RES]]
;
%mul = mul nsw i32 %x, 5
%res = ashr i32 %mul, 2
call void @use(i32 %mul)
ret i32 %res
}

define i32 @ashr_mul_times_5_div_4_exact_2(i32 %x) {
; CHECK-LABEL: @ashr_mul_times_5_div_4_exact_2(
; CHECK-NEXT: [[TMP1:%.*]] = ashr exact i32 [[X:%.*]], 2
; CHECK-NEXT: [[ASHR:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[ASHR]]
;
%mul = mul nsw i32 %x, 5
%ashr = ashr exact i32 %mul, 2
ret i32 %ashr
}

declare void @use(i32)
19 changes: 16 additions & 3 deletions llvm/test/Transforms/InstCombine/lshr.ll
View file
Original file line number Diff line number Diff line change
Expand Up @@ -628,19 +628,32 @@ define i32 @mul_splat_fold_wrong_lshr_const(i32 %x) {
ret i32 %t
}

; Negative test
; Negative test (but simplifies into a different transform)

define i32 @mul_splat_fold_no_nuw(i32 %x) {
; CHECK-LABEL: @mul_splat_fold_no_nuw(
; CHECK-NEXT: [[M:%.*]] = mul nsw i32 [[X:%.*]], 65537
; CHECK-NEXT: [[T:%.*]] = lshr i32 [[M]], 16
; CHECK-NEXT: [[TMP1:%.*]] = lshr i32 [[X:%.*]], 16
; CHECK-NEXT: [[T:%.*]] = add nsw i32 [[TMP1]], [[X]]
; CHECK-NEXT: ret i32 [[T]]
;
%m = mul nsw i32 %x, 65537
%t = lshr i32 %m, 16
ret i32 %t
}

; Negative test

define i32 @mul_splat_fold_no_flags(i32 %x) {
; CHECK-LABEL: @mul_splat_fold_no_flags(
; CHECK-NEXT: [[M:%.*]] = mul i32 [[X:%.*]], 65537
; CHECK-NEXT: [[T:%.*]] = lshr i32 [[M]], 16
; CHECK-NEXT: ret i32 [[T]]
;
%m = mul i32 %x, 65537
%t = lshr i32 %m, 16
ret i32 %t
}

; Negative test (but simplifies before we reach the mul_splat transform)- need more than 2 bits

define i2 @mul_splat_fold_too_narrow(i2 %x) {
Expand Down

0 comments on commit 718ba5a

Please sign in to comment.