[InstCombine] Fold (X << Z) / (X * Y) -> (1 << Z) / Y

[Chenbing.Zheng]

Summary

Alive2: https://alive2.llvm.org/ce/z/CBJLeP

Note that both shl and mul require nuw and already
have negtive test cases in div-shift.ll for it.

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

Can you preserve NUW and exact?

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

556–629

Vector test? Tests with flag preservation?

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

Move the !IsSigned check outside to avoid the cost of match calls?

[Chenbing.Zheng]

1. rebase after add more tests
2. add oneuse check
3. keep exact flag

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

I preserve exact, but it seems that NUW is unnecessary when Op0 of SHL is 1 ?

[Chenbing.Zheng]

ping...

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

You're right that we can infer nuw for 1 << x, but based on your tests it doesn't look like we actually do this...

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

Yer, I think 'shl nuw 1 << x' is equivalent to 'shl 1 << x', so I tend to don't keep 'NUW'. Or what is the benefit of keeping 'NUW' ?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1066

It could enable some other transform. Some tests actually improve with that:
d4493dd1ed58

So assuming that patch is good, it won't matter much now if it isn't included here, but it's still more efficient to create the 'nuw' directly in this patch instead of relying on another fold to do it.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1061

Do you plan to handle the signed case too?

I'm not sure exactly what no-wrap is required, but there is some combination that works, so add a TODO comment:
https://alive2.llvm.org/ce/z/WrXmth

[Chenbing.Zheng]

add nuw for 1<<Z

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1062

Would this still be a good transform even if the shift has another use (is there a test for that pattern)?
If we are trading a multiply for a shift, it's probably beneficial.

[Chenbing.Zheng]

address comment

[spatel]

LGTM

[Chenbing.Zheng]

https://reviews.llvm.org/rG1f84e72b7ba0a563e4abf2320997887bf8ae3f7c

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (match(Op1, m_c_Mul(m_Specific(Op0), m_Value(Y)))) {
     bool HasNUW = cast<OverflowingBinaryOperator>(Op1)->hasNoUnsignedWrap();
     if ((IsSigned && HasNSW) || (!IsSigned && HasNUW)) {
       replaceOperand(I, 0, ConstantInt::get(Ty, 1));
       replaceOperand(I, 1, Y);
       return &I;
     }
   }
 
+  // (X << Z) / (X * Y) -> (1 << Z) / Y
+  // TODO: Handle sdiv.
+  if (!IsSigned && Op1->hasOneUse() &&
+      match(Op0, m_NUWShl(m_Value(X), m_Value(Z))) &&
+      match(Op1, m_c_Mul(m_Specific(X), m_Value(Y))))
+    if (cast<OverflowingBinaryOperator>(Op1)->hasNoUnsignedWrap()) {
+      Instruction *NewDiv = BinaryOperator::CreateUDiv(
+          Builder.CreateShl(ConstantInt::get(Ty, 1), Z, "", /*NUW*/ true), Y);
+      NewDiv->setIsExact(I.isExact());
+      return NewDiv;
+    }
+
   if (Instruction *R = foldIDivShl(I, Builder))
     return R;
 
   // With the appropriate no-wrap constraint, remove a multiply by the divisor
   // after peeking through another divide:
   // ((Op1 * X) / Y) / Op1 --> X / Y
   if (match(Op0, m_BinOp(I.getOpcode(), m_c_Mul(m_Specific(Op1), m_Value(X)),
                          m_Value(Y)))) {

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

;
   %m1 = mul nuw i8 %x, %y
   call void @use(i8 %m1)
   %m2 = shl nuw i8 %x, %z
   call void @use(i8 %m2)
   %d = udiv i8 %m1, %m2
   ret i8 %d
 }
 
-; TODO: This can fold to (1<<z) / y
+; (X << Z) / (X * Y) -> (1 << Z) / Y
 
 define i5 @udiv_shl_mul_nuw(i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @udiv_shl_mul_nuw(
-; CHECK-NEXT:    [[M1:%.*]] = shl nuw i5 [[X:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[M2:%.*]] = mul nuw i5 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[D:%.*]] = udiv i5 [[M1]], [[M2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i5 1, [[Z:%.*]]
+; CHECK-NEXT:    [[D:%.*]] = udiv i5 [[TMP1]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i5 [[D]]
 ;
   %m1 = shl nuw i5 %x, %z
   %m2 = mul nuw i5 %x, %y
   %d = udiv i5 %m1, %m2
   ret i5 %d
 }
 
 define i5 @udiv_shl_mul_nuw_swap(i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @udiv_shl_mul_nuw_swap(
-; CHECK-NEXT:    [[M1:%.*]] = shl nuw i5 [[X:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[M2:%.*]] = mul nuw i5 [[Y:%.*]], [[X]]
-; CHECK-NEXT:    [[D:%.*]] = udiv i5 [[M1]], [[M2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i5 1, [[Z:%.*]]
+; CHECK-NEXT:    [[D:%.*]] = udiv i5 [[TMP1]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i5 [[D]]
 ;
   %m1 = shl nuw i5 %x, %z
   %m2 = mul nuw i5 %y, %x
   %d = udiv i5 %m1, %m2
   ret i5 %d
 }
 
 define i5 @udiv_shl_mul_nuw_exact(i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @udiv_shl_mul_nuw_exact(
-; CHECK-NEXT:    [[M1:%.*]] = shl nuw i5 [[X:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[M2:%.*]] = mul nuw i5 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[D:%.*]] = udiv exact i5 [[M1]], [[M2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i5 1, [[Z:%.*]]
+; CHECK-NEXT:    [[D:%.*]] = udiv exact i5 [[TMP1]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i5 [[D]]
 ;
   %m1 = shl nuw i5 %x, %z
   %m2 = mul nuw i5 %x, %y
   %d = udiv exact i5 %m1, %m2
   ret i5 %d
 }
 
 define <2 x i4> @udiv_shl_mul_nuw_vec(<2 x i4> %x, <2 x i4> %y, <2 x i4> %z) {
 ; CHECK-LABEL: @udiv_shl_mul_nuw_vec(
-; CHECK-NEXT:    [[M1:%.*]] = shl nuw <2 x i4> [[X:%.*]], [[Z:%.*]]
-; CHECK-NEXT:    [[M2:%.*]] = mul nuw <2 x i4> [[Y:%.*]], [[X]]
-; CHECK-NEXT:    [[D:%.*]] = udiv <2 x i4> [[M1]], [[M2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw <2 x i4> <i4 1, i4 1>, [[Z:%.*]]
+; CHECK-NEXT:    [[D:%.*]] = udiv <2 x i4> [[TMP1]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i4> [[D]]
 ;
   %m1 = shl nuw <2 x i4> %x, %z
   %m2 = mul nuw <2 x i4> %y, %x
   %d = udiv <2 x i4> %m1, %m2
   ret <2 x i4> %d
 }
 
+define i8 @udiv_shl_mul_nuw_extra_use_of_shl(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @udiv_shl_mul_nuw_extra_use_of_shl(
+; CHECK-NEXT:    [[M1:%.*]] = shl nuw i8 [[X:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    call void @use(i8 [[M1]])
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i8 1, [[Z]]
+; CHECK-NEXT:    [[D:%.*]] = udiv i8 [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[D]]
+;
+  %m1 = shl nuw i8 %x, %z
+  call void @use(i8 %m1)
+  %m2 = mul nuw i8 %y, %x
+  %d = udiv i8 %m1, %m2
+  ret i8 %d
+}
+
 ; negative test - extra use
 
+define i8 @udiv_shl_mul_nuw_extra_use_of_mul(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @udiv_shl_mul_nuw_extra_use_of_mul(
+; CHECK-NEXT:    [[M1:%.*]] = shl nuw i8 [[X:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[M2:%.*]] = mul nuw i8 [[Y:%.*]], [[X]]
+; CHECK-NEXT:    call void @use(i8 [[M2]])
+; CHECK-NEXT:    [[D:%.*]] = udiv i8 [[M1]], [[M2]]
+; CHECK-NEXT:    ret i8 [[D]]
+;
+  %m1 = shl nuw i8 %x, %z
+  %m2 = mul nuw i8 %y, %x
+  call void @use(i8 %m2)
+  %d = udiv i8 %m1, %m2
+  ret i8 %d
+}
+
 define i8 @udiv_shl_mul_nuw_extra_use(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_shl_mul_nuw_extra_use(
 ; CHECK-NEXT:    [[M1:%.*]] = shl nuw i8 [[X:%.*]], [[Z:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[M1]])
 ; CHECK-NEXT:    [[M2:%.*]] = mul nuw i8 [[Y:%.*]], [[X]]
 ; CHECK-NEXT:    call void @use(i8 [[M2]])
 ; CHECK-NEXT:    [[D:%.*]] = udiv i8 [[M1]], [[M2]]
 ; CHECK-NEXT:    ret i8 [[D]]
 ;
   %m1 = shl nuw i8 %x, %z
   call void @use(i8 %m1)
   %m2 = mul nuw i8 %y, %x
   call void @use(i8 %m2)
   %d = udiv i8 %m1, %m2
   ret i8 %d
 }
 
 ; negative test - sdiv
 
 define i5 @sdiv_shl_mul_nuw(i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @sdiv_shl_mul_nuw(
 ; CHECK-NEXT:    [[M1:%.*]] = shl nuw i5 [[X:%.*]], [[Z:%.*]]
 ; CHECK-NEXT:    [[M2:%.*]] = mul nuw i5 [[X]], [[Y:%.*]]
 ; CHECK-NEXT:    [[D:%.*]] = sdiv i5 [[M1]], [[M2]]
 ; CHECK-NEXT:    ret i5 [[D]]