[llvm] 9cff471 - [InstCombine] fold udiv with common factor

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2022-10-10T08:12:06-04:00
New Revision: 9cff4711ac72576014a08f0cc0c773fa7bcc5557

URL: https://github.com/llvm/llvm-project/commit/9cff4711ac72576014a08f0cc0c773fa7bcc5557
DIFF: https://github.com/llvm/llvm-project/commit/9cff4711ac72576014a08f0cc0c773fa7bcc5557.diff

LOG: [InstCombine] fold udiv with common factor

((X *nuw Y) >> Z) / X --> Y >> Z

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

This is similar to 6b869be8100d / 8da2fa856f1b, but I have
not found a signed equivalent, so it's just an unsigned
match for now.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
    llvm/test/Transforms/InstCombine/div-shift.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index d9fcd436b00e0..0bea26ac57cd0 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -1157,6 +1157,16 @@ Instruction *InstCombinerImpl::visitUDiv(BinaryOperator &I) {
       return BinaryOperator::CreateUDiv(A, X);
   }
 
+  // Look through a right-shift to find the common factor:
+  // ((Op1 *nuw A) >> B) / Op1 --> A >> B
+  if (match(Op0, m_LShr(m_NUWMul(m_Specific(Op1), m_Value(A)), m_Value(B))) ||
+      match(Op0, m_LShr(m_NUWMul(m_Value(A), m_Specific(Op1)), m_Value(B)))) {
+    Instruction *Lshr = BinaryOperator::CreateLShr(A, B);
+    if (I.isExact() && cast<PossiblyExactOperator>(Op0)->isExact())
+      Lshr->setIsExact();
+    return Lshr;
+  }
+
   // Op1 udiv Op2 -> Op1 lshr log2(Op2), if log2() folds away.
   if (takeLog2(Builder, Op1, /*Depth*/0, /*DoFold*/false)) {
     Value *Res = takeLog2(Builder, Op1, /*Depth*/0, /*DoFold*/true);

diff  --git a/llvm/test/Transforms/InstCombine/div-shift.ll b/llvm/test/Transforms/InstCombine/div-shift.ll
index 101cbf11ea6f0..c266fa76247b1 100644
--- a/llvm/test/Transforms/InstCombine/div-shift.ll
+++ b/llvm/test/Transforms/InstCombine/div-shift.ll
@@ -619,11 +619,11 @@ define i8 @udiv_shl_nuw_use(i8 %x, i8 %y, i8 %z) {
   ret i8 %d
 }
 
+; ((X * Y) >> Z) / X --> Y >> Z
+
 define i8 @udiv_lshr_mul_nuw(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw(
-; CHECK-NEXT:    [[M:%.*]] = mul nuw i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[S:%.*]] = lshr i8 [[M]], [[Z:%.*]]
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 [[S]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = lshr i8 [[Y:%.*]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[DIV]]
 ;
   %m = mul nuw i8 %x, %y
@@ -632,11 +632,11 @@ define i8 @udiv_lshr_mul_nuw(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; ((Y * X) >> Z) / X --> Y >> Z
+
 define <2 x i4> @udiv_lshr_mul_nuw_exact_commute1(<2 x i4> %x, <2 x i4> %y, <2 x i4> %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw_exact_commute1(
-; CHECK-NEXT:    [[M:%.*]] = mul nuw <2 x i4> [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[S:%.*]] = lshr exact <2 x i4> [[M]], [[Z:%.*]]
-; CHECK-NEXT:    [[DIV:%.*]] = udiv exact <2 x i4> [[S]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = lshr exact <2 x i4> [[Y:%.*]], [[Z:%.*]]
 ; CHECK-NEXT:    ret <2 x i4> [[DIV]]
 ;
   %m = mul nuw <2 x i4> %y, %x
@@ -645,6 +645,8 @@ define <2 x i4> @udiv_lshr_mul_nuw_exact_commute1(<2 x i4> %x, <2 x i4> %y, <2 x
   ret <2 x i4> %div
 }
 
+; negative test - mul is shifted amount, not shifted value
+
 define i8 @udiv_lshr_mul_nuw_commute2(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw_commute2(
 ; CHECK-NEXT:    [[M:%.*]] = mul nuw i8 [[Y:%.*]], [[X:%.*]]
@@ -658,12 +660,13 @@ define i8 @udiv_lshr_mul_nuw_commute2(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; extra uses are ok
+
 define i8 @udiv_lshr_mul_nuw_use1(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw_use1(
 ; CHECK-NEXT:    [[M:%.*]] = mul nuw i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[M]])
-; CHECK-NEXT:    [[S:%.*]] = lshr i8 [[M]], [[Z:%.*]]
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 [[S]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = lshr i8 [[Y]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[DIV]]
 ;
   %m = mul nuw i8 %x, %y
@@ -673,12 +676,14 @@ define i8 @udiv_lshr_mul_nuw_use1(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; extra uses are ok
+
 define i8 @udiv_lshr_mul_nuw_use2(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw_use2(
 ; CHECK-NEXT:    [[M:%.*]] = mul nuw i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[S:%.*]] = lshr i8 [[M]], [[Z:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[S]])
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 [[S]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = lshr i8 [[Y]], [[Z]]
 ; CHECK-NEXT:    ret i8 [[DIV]]
 ;
   %m = mul nuw i8 %x, %y
@@ -688,13 +693,15 @@ define i8 @udiv_lshr_mul_nuw_use2(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; extra uses are ok
+
 define i8 @udiv_lshr_mul_nuw_use3(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nuw_use3(
 ; CHECK-NEXT:    [[M:%.*]] = mul nuw i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[M]])
 ; CHECK-NEXT:    [[S:%.*]] = lshr i8 [[M]], [[Z:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[S]])
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 [[S]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = lshr i8 [[Y]], [[Z]]
 ; CHECK-NEXT:    ret i8 [[DIV]]
 ;
   %m = mul nuw i8 %x, %y
@@ -705,6 +712,8 @@ define i8 @udiv_lshr_mul_nuw_use3(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; negative test - must have nuw
+
 define i8 @udiv_lshr_mul_nsw(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_lshr_mul_nsw(
 ; CHECK-NEXT:    [[M:%.*]] = mul nsw i8 [[X:%.*]], [[Y:%.*]]
@@ -718,6 +727,8 @@ define i8 @udiv_lshr_mul_nsw(i8 %x, i8 %y, i8 %z) {
   ret i8 %div
 }
 
+; negative test - doesn't fold with signed div
+
 define i8 @sdiv_lshr_mul_nsw(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @sdiv_lshr_mul_nsw(
 ; CHECK-NEXT:    [[M:%.*]] = mul nsw i8 [[X:%.*]], [[Y:%.*]]