[InstCombine] Fold mul nuw+lshr to a single multiplication when the latter is a factor

[bcl5980]

Summary

if c is divisible by (1 << ShAmtC), we can fold this pattern:
lshr (mul nuw x, c), ShAmtC -> mul nuw x, (c >> ShAmtC)

https://alive2.llvm.org/ce/z/ox4wAt

Fix https://github.com/llvm/llvm-project/issues/54824

[xbolva00]

Comment should be more clear

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

mul nuw, c??

1173

qc is what?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

I'm sorry, should be :
lshr exact (mul nuw x, c * (1 << ShAmtC)), ShAmtC -> mul nuw, c
I will fix the comment later

[bcl5980]

fix typo

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

mul nuw x, c?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

yeah

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

c * (1 << ShAmtC) is (c << ShAmtC) isn't it?

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1166–1167

This comment needs to move inside the first.

[lebedev.ri]

Please can you post the link to the general proof?

[bcl5980]

In D123453#3441088, @lebedev.ri wrote:

Please can you post the link to the general proof?

I'm sorry but can you help to explain what the general proof is?

[bcl5980]

Fix comment
Update alive2 link

[lebedev.ri]

In D123453#3441114, @bcl5980 wrote:

In D123453#3441088, @lebedev.ri wrote:

Please can you post the link to the general proof?

I'm sorry but can you help to explain what the general proof is?

The one you linked has hardcoded constants, yet the constants aren't hardcoded in the transform itself.

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1177

Once again:

mul nuw x, c

[bcl5980]

Fix comment
Update general proof link

[bcl5980]

In D123453#3441118, @lebedev.ri wrote:

In D123453#3441114, @bcl5980 wrote:

In D123453#3441088, @lebedev.ri wrote:

Please can you post the link to the general proof?

I'm sorry but can you help to explain what the general proof is?

The one you linked has hardcoded constants, yet the constants aren't hardcoded in the transform itself.

Thanks for the mentioning. I'll pay attention next time for constants.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1177

Sorry for that. Miss again

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

1. No one-use check needed, only a single instruction is produced.
2. I believe, while that lshr is obviously exact, your own proof shows that there is no need to check that it is marked as such?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

If there is no one-use, we may generate 2 mul instead of mul+shr

define i64 @lshr_mul_negative_oneuse(i64 %0) {
; CHECK-LABEL: @lshr_mul_negative_oneuse(
; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw i64 [[TMP0:%.*]], 52
; CHECK-NEXT:    call void @use(i64 [[TMP2]])
; CHECK-NEXT:    [[TMP3:%.*]] = mul nuw i64 [[TMP0]], 13
; CHECK-NEXT:    ret i64 [[TMP3]]
;
  %2 = mul nuw i64 %0, 52
  call void @use(i64 %2)
  %3 = lshr i64 %2, 2
  ret i64 %3
}

We need a condition to make sure ShAmtC is divisible by NewMulC. The proof use shl that is always sure. But in the code we still need to check eact flag.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

We need a condition to make sure ShAmtC is divisible by NewMulC. The proof use shl that is always sure. But in the code we still need to check eact flag.

Should be "make sure MulC is divisible by NewMulC"

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

If there is no one-use, we may generate 2 mul instead of mul+shr

We only create only a single mul here, do we not?

We need a condition to make sure ShAmtC is divisible by NewMulC. The proof use shl that is always sure. But in the code we still need to check eact flag.

Can you show the counter-proof that shows that not checking for exact is incorrect?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

If there is no one-use, we may generate 2 mul instead of mul+shr

We only create only a single mul here, do we not?

Yeah, we create only a single mul, but most of time mul should be heavier than lshr, am I right?

We need a condition to make sure ShAmtC is divisible by NewMulC. The proof use shl that is always sure. But in the code we still need to check eact flag.

Can you show the counter-proof that shows that not checking for exact is incorrect?

I don't know how to counter-proof with alive2 but this is the negative case on my machine after remove exact:

define i64 @lshr_mul_negative_noexact(i64 %0) {
; CHECK-LABEL: @lshr_mul_negative_noexact(
; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw i64 [[TMP0:%.*]], 13
; CHECK-NEXT:    ret i64 [[TMP2]]
;
  %2 = mul nuw i64 %0, 53
  %3 = lshr i64 %2, 2
  ret i64 %3
}

[lebedev.ri]

I see. Then the proof is wrong.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

If there is no one-use, we may generate 2 mul instead of mul+shr

We only create only a single mul here, do we not?

declare i64 @use(i64, i64)
define i64 @lshr_mul_negative_oneuse(i64 %0) {
; CHECK-LABEL: @lshr_mul_negative_oneuse(
; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw i64 [[TMP0:%.*]], 52
; CHECK-NEXT:    [[TMP3:%.*]] = mul nuw i64 [[TMP0]], 13
; CHECK-NEXT:    [[TMP4:%.*]] = call i64 @use(i64 [[TMP2]], i64 [[TMP3]])
; CHECK-NEXT:    ret i64 [[TMP4]]
;
  %2 = mul nuw i64 %0, 52
  %3 = lshr i64 %2, 2
  %4 = call i64 @use(i64 %2, i64 %3)
  ret i64 %4
}

This case should more clear why we need one use I think.

[bcl5980]

Fix test error

[bcl5980]

In D123453#3441148, @lebedev.ri wrote:

I see. Then the proof is wrong.

I'm sorry I really don't know how to proof by alive2. Can you teach me how to proof it?

[bcl5980]

Fix comment

[lebedev.ri]

In D123453#3441197, @bcl5980 wrote:

In D123453#3441148, @lebedev.ri wrote:

I see. Then the proof is wrong.

I'm sorry I really don't know how to proof by alive2. Can you teach me how to proof it?

You need to either adjust the fold to do what the proof says, or write another proof for the change in question.

Abstracting a bit, is there a generalization of this pattern where the shift amounts don't match?

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

The number of instructions didn't increase, so we're fine.

[lebedev.ri]

For example, what about this:
https://alive2.llvm.org/ce/z/aov-GB
https://alive2.llvm.org/ce/z/ox4wAt (no need for exact, only nuw)

[bcl5980]

In D123453#3441216, @lebedev.ri wrote:

For example, what about this:
https://alive2.llvm.org/ce/z/aov-GB
https://alive2.llvm.org/ce/z/ox4wAt (no need for exact, only nuw)

%t0 = lshr i8 %C1, %C2
%t1 = shl i8 %t0, %C2
%precond = icmp eq i8 %t1, %C1
call void @llvm.assume(i1 %precond)

This is the proof of exact. We can't remove exact check in the code. Or we need to do something like:

if (Op0->hasOneUse()) {
  APInt NewMulC = MulC->lshr(ShAmtC);
  if (MulC->eq(NewMulC.shl(ShAmtC)))
    return BinaryOperator::CreateNUWMul(X, ConstantInt::get(Ty, NewMulC));
}

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1178

The number of instructions didn't increase, so we're fine.

I still insist one-use, mul is slower than lshr on most targets.

[craig.topper]

In D123453#3441326, @bcl5980 wrote:

In D123453#3441216, @lebedev.ri wrote:

For example, what about this:
https://alive2.llvm.org/ce/z/aov-GB
https://alive2.llvm.org/ce/z/ox4wAt (no need for exact, only nuw)

%t0 = lshr i8 %C1, %C2
%t1 = shl i8 %t0, %C2
%precond = icmp eq i8 %t1, %C1
call void @llvm.assume(i1 %precond)

This is the proof of exact. We can't remove exact check in the code. Or we need to do something like:

if (Op0->hasOneUse()) {
  APInt NewMulC = MulC->lshr(ShAmtC);
  if (MulC->eq(NewMulC.shl(ShAmtC)))
    return BinaryOperator::CreateNUWMul(X, ConstantInt::get(Ty, NewMulC));
}

Couldn't the exact be set because the LHS of the multiply is known to have trailing zero bits and have nothing to do with the trailing zeros of the constant on the RHS. Would the transform still be valid in that case?

[bcl5980]

In D123453#3441329, @craig.topper wrote:

In D123453#3441326, @bcl5980 wrote:

In D123453#3441216, @lebedev.ri wrote:

For example, what about this:
https://alive2.llvm.org/ce/z/aov-GB
https://alive2.llvm.org/ce/z/ox4wAt (no need for exact, only nuw)

%t0 = lshr i8 %C1, %C2
%t1 = shl i8 %t0, %C2
%precond = icmp eq i8 %t1, %C1
call void @llvm.assume(i1 %precond)

This is the proof of exact. We can't remove exact check in the code. Or we need to do something like:

if (Op0->hasOneUse()) {
  APInt NewMulC = MulC->lshr(ShAmtC);
  if (MulC->eq(NewMulC.shl(ShAmtC)))
    return BinaryOperator::CreateNUWMul(X, ConstantInt::get(Ty, NewMulC));
}

Couldn't the exact be set because the LHS of the multiply is known to have trailing zero bits and have nothing to do with the trailing zeros of the constant on the RHS. Would the transform still be valid in that case?

Thanks for the finding. That's what I haven't expect this case. So we should use the check if (MulC->eq(NewMulC.shl(ShAmtC))) to avoid mul's LHS involve exact.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

Not related to this patch, but this existing transform is incorrect if MulC happens to be 2 and BitWidth is 2 and ShAmtC is 1. MulC-1 would be 1 which is a power 2 but it doesn't have 2 bits set like the transform expects. Hopefully we would always canonicalize the Mul to a shl first so this transform won't fire for that case, but I'm not sure. @spatel @lebedev.ri

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1173

Good catch. I'm not sure how to write a test for it either, but I added a "BitWidth > 2" clause:
1206a18d417a

[bcl5980]

rebase

[bcl5980]

fix build error

[bcl5980]

@lebedev.ri @craig.topper
What can I do now to continue the patch?

[bcl5980]

Ping.

[spatel]

1. Please pre-commit the tests with baseline CHECKs. Add one more positive test with vector type, so we know that works correctly -- at least for splat (uniform) constants. A negative test with 'nsw' would also be good. What happens if we have both 'nsw' and 'nuw'?

2. There's an existing fold for: (X * C2) << C1 --> X * (C2 << C1)

...and it does not check for one-use. For consistency, we probably don't want the one-use restriction here either. If there's already a multiply in the pattern before this transform, another one is probably fine? The backend could theoretically decompose it back to shift (but I don't think we have that transform currently).

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1179–1180

Replace 'c' with 'MulC' in this comment to make it clearer.

[bcl5980]

void test(unsigned* pa, unsigned* pb, unsigned short i)
{
    unsigned a = i * 100;
    unsigned b = a >> 2;
    *pa = a;
    *pb = b;
}

I write a test to verify backend can optimize the pattern without one use.
This is the AArch64 baseline with one-use:

// %bb.0:                               // %entry
	and	w8, w2, #0xffff
	mov	w9, #100
	mul	w8, w8, w9
	lsr	w9, w8, #2
	str	w8, [x0]
	str	w9, [x1]
	ret
                                        // -- End function

This is the AArch64 result we remove one-use:

"?test@@YAXPEAI0G@Z":                   // @"?test@@YAXPEAI0G@Z"
// %bb.0:                               // %entry
	mov	w8, #100
	and	w9, w2, #0xffff
	mov	w10, #25
	mul	w8, w9, w8
	mul	w9, w9, w10
	str	w8, [x0]
	str	w9, [x1]
	ret
                                        // -- End function

This is the X86 baseline with one-use:

"?test@@YAXPEAI0G@Z":                   # @"?test@@YAXPEAI0G@Z"
# %bb.0:                                # %entry
	movzwl	%r8w, %eax
	imull	$100, %eax, %eax
	movl	%eax, (%rcx)
	shrl	$2, %eax
	movl	%eax, (%rdx)
	retq
                                        # -- End function

This is the X86 result we remove one-use:

"?test@@YAXPEAI0G@Z":                   # @"?test@@YAXPEAI0G@Z"
# %bb.0:                                # %entry
	movzwl	%r8w, %eax
	imull	$100, %eax, %r8d
	leal	(%rax,%rax,4), %eax
	leal	(%rax,%rax,4), %eax
	movl	%r8d, (%rcx)
	movl	%eax, (%rdx)
	retq
                                        # -- End function

Backend is not easy to figure out the case as we need to loop all use of mul's operand 0 to find the candidate.

In D123453#3459579, @spatel wrote:

2. There's an existing fold for: (X * C2) << C1 --> X * (C2 << C1)

...and it does not check for one-use. For consistency, we probably don't want the one-use restriction here either. If there's already a multiply in the pattern before this transform, another one is probably fine? The backend could theoretically decompose it back to shift (but I don't think we have that transform currently).

Can we add one-use for this pattern also? https://godbolt.org/z/x3bo7q54j

[bcl5980]

update comments and rebase

[spatel]

In D123453#3461148, @bcl5980 wrote:

Backend is not easy to figure out the case as we need to loop all use of mul's operand 0 to find the candidate.

Ok - I agree that the posted examples don't look like wins. There is no quick fix for the backend, so we can leave the one-use check in for now. Please add a code comment though to explain why have the one-use limitation.

If you want to make the other transform (for "shl") consistent with this one by adding a one-use check to it too, that can be a follow-up patch. That transform has been around for a long time, so it is possible that someone may notice an independent performance difference from that change.

llvm/test/Transforms/InstCombine/shift-logic.ll

275

If we do not propagate the 'nsw' on this, we will not be able to recover it in general:
https://alive2.llvm.org/ce/z/tbckkt

[bcl5980]

keep nsw flag

[spatel]

LGTM - see inline remark for minor edit to a code comment.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1177

I would phrase this differently:

// 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 instruction.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op0, m_OneUse(m_Trunc(m_Instruction(TruncSrc)))) &&
         Value *Trunc = Builder.CreateTrunc(SumShift, Ty, I.getName());
 
         // If the first shift does not cover the number of bits truncated, then
         // we require a mask to get rid of high bits in the result.
         APInt MaskC = APInt::getAllOnes(BitWidth).lshr(ShAmtC);
         return BinaryOperator::CreateAnd(Trunc, ConstantInt::get(Ty, MaskC));
       }
     }
 
+    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?
-    const APInt *MulC;
-    if (match(Op0, m_NUWMul(m_Value(X), m_APInt(MulC))) &&
-        BitWidth > 2 && ShAmtC * 2 == BitWidth && (*MulC - 1).isPowerOf2() &&
+      if (BitWidth > 2 && ShAmtC * 2 == BitWidth && (*MulC - 1).isPowerOf2() &&
           MulC->logBase2() == ShAmtC)
         return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, *MulC - 2));
 
+      // Check one-use here to make sure we can save the shift instruction,
+      // avoid replace shift to mul as shift is cheaper than mul in general.
+      if (Op0->hasOneUse()) {
+        APInt NewMulC = MulC->lshr(ShAmtC);
+        // if c is divisible by (1 << ShAmtC):
+        // lshr (mul nuw x, MulC), ShAmtC -> mul nuw x, (MulC >> ShAmtC)
+        if (MulC->eq(NewMulC.shl(ShAmtC))) {
+          auto *NewMul =
+              BinaryOperator::CreateNUWMul(X, ConstantInt::get(Ty, NewMulC));
+          BinaryOperator *OrigMul = cast<BinaryOperator>(Op0);
+          NewMul->setHasNoSignedWrap(OrigMul->hasNoSignedWrap());
+          return NewMul;
+        }
+      }
+    }
+
     // Try to narrow a bswap:
     // (bswap (zext X)) >> C --> zext (bswap X >> C')
     // In the case where the shift amount equals the bitwidth difference, the
     // shift is eliminated.
     if (match(Op0, m_OneUse(m_Intrinsic<Intrinsic::bswap>(
                        m_OneUse(m_ZExt(m_Value(X))))))) {
       // TODO: If the shift amount is less than the zext, we could shift left.
       unsigned SrcWidth = X->getType()->getScalarSizeInBits();

llvm/test/Transforms/InstCombine/shift-logic.ll

;
   %sh1 = ashr exact i32 %x, 16
   %t0 = xor i32 %sh1, shl (i32 ptrtoint (i32* @g to i32), i32 16)
   %t27 = ashr exact i32 %t0, 16
   ret i32 %t27
 }
 
 define i64 @lshr_mul(i64 %0) {
 ; CHECK-LABEL: @lshr_mul(
-; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw i64 [[TMP0:%.*]], 52
-; CHECK-NEXT:    [[TMP3:%.*]] = lshr exact i64 [[TMP2]], 2
-; CHECK-NEXT:    ret i64 [[TMP3]]
+; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw i64 [[TMP0:%.*]], 13
+; CHECK-NEXT:    ret i64 [[TMP2]]
 ;
   %2 = mul nuw i64 %0, 52
   %3 = lshr i64 %2, 2
   ret i64 %3
 }
 
 define i64 @lshr_mul_nuw_nsw(i64 %0) {
 ; CHECK-LABEL: @lshr_mul_nuw_nsw(
-; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw nsw i64 [[TMP0:%.*]], 52
-; CHECK-NEXT:    [[TMP3:%.*]] = lshr exact i64 [[TMP2]], 2
-; CHECK-NEXT:    ret i64 [[TMP3]]
+; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw nsw i64 [[TMP0:%.*]], 13
+; CHECK-NEXT:    ret i64 [[TMP2]]
 ;
   %2 = mul nuw nsw i64 %0, 52
   %3 = lshr i64 %2, 2
   ret i64 %3
 }
 
 define <4 x i32> @lshr_mul_vector(<4 x i32> %0) {
 ; CHECK-LABEL: @lshr_mul_vector(
-; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw <4 x i32> [[TMP0:%.*]], <i32 52, i32 52, i32 52, i32 52>
-; CHECK-NEXT:    [[TMP3:%.*]] = lshr exact <4 x i32> [[TMP2]], <i32 2, i32 2, i32 2, i32 2>
-; CHECK-NEXT:    ret <4 x i32> [[TMP3]]
+; CHECK-NEXT:    [[TMP2:%.*]] = mul nuw <4 x i32> [[TMP0:%.*]], <i32 13, i32 13, i32 13, i32 13>
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
 ;
   %2 = mul nuw <4 x i32> %0, <i32 52, i32 52, i32 52, i32 52>
   %3 = lshr <4 x i32> %2, <i32 2, i32 2, i32 2, i32 2>
   ret <4 x i32> %3
 }
 
 define i64 @lshr_mul_negative_noexact(i64 %0) {
 ; CHECK-LABEL: @lshr_mul_negative_noexact(

 ; CHECK-NEXT:    [[TMP2:%.*]] = mul i64 [[TMP0:%.*]], 52
 ; CHECK-NEXT:    [[TMP3:%.*]] = lshr exact i64 [[TMP2]], 2
 ; CHECK-NEXT:    ret i64 [[TMP3]]
 ;
   %2 = mul i64 %0, 52
   %3 = lshr i64 %2, 2
   ret i64 %3
 }
+
+define i64 @lshr_mul_negative_nsw(i64 %0) {
+; CHECK-LABEL: @lshr_mul_negative_nsw(
+; CHECK-NEXT:    [[TMP2:%.*]] = mul nsw i64 [[TMP0:%.*]], 52
+; CHECK-NEXT:    [[TMP3:%.*]] = lshr exact i64 [[TMP2]], 2
+; CHECK-NEXT:    ret i64 [[TMP3]]
+;
+  %2 = mul nsw i64 %0, 52
+  %3 = lshr i64 %2, 2
+  ret i64 %3
+}