[llvm] afa3d58 - [InstCombine] Fold `(mul (div exact X, C0), C1)` -> `(div exact X, C0/C1)`

[Noah Goldstein via llvm-commits]

Author: Noah Goldstein
Date: 2024-06-28T16:52:03+08:00
New Revision: afa3d58ee2224f1037116f4cab44a23bd232a416

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

LOG: [InstCombine] Fold `(mul (div exact X, C0), C1)` -> `(div exact X, C0/C1)`

We can do this if `C0 % C1 == 0` and if we avoid UB in the signed
case.

Proofs: https://alive2.llvm.org/ce/z/HHWHDg

Closes #96915

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
    llvm/test/Transforms/InstCombine/exact.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 0824ca02c1289..3a2620db1f042 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -365,6 +365,28 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
       return BinaryOperator::CreateMul(NegOp0, X);
   }
 
+  if (Op0->hasOneUse()) {
+    // (mul (div exact X, C0), C1)
+    //    -> (div exact X, C0 / C1)
+    // iff C0 % C1 == 0 and X / (C0 / C1) doesn't create UB.
+    const APInt *C1;
+    auto UDivCheck = [&C1](const APInt &C) { return C.urem(*C1).isZero(); };
+    auto SDivCheck = [&C1](const APInt &C) {
+      APInt Quot, Rem;
+      APInt::sdivrem(C, *C1, Quot, Rem);
+      return Rem.isZero() && !Quot.isAllOnes();
+    };
+    if (match(Op1, m_APInt(C1)) &&
+        (match(Op0, m_Exact(m_UDiv(m_Value(X), m_CheckedInt(UDivCheck)))) ||
+         match(Op0, m_Exact(m_SDiv(m_Value(X), m_CheckedInt(SDivCheck)))))) {
+      auto BOpc = cast<BinaryOperator>(Op0)->getOpcode();
+      return BinaryOperator::CreateExact(
+          BOpc, X,
+          Builder.CreateBinOp(BOpc, cast<BinaryOperator>(Op0)->getOperand(1),
+                              Op1));
+    }
+  }
+
   // (X / Y) *  Y = X - (X % Y)
   // (X / Y) * -Y = (X % Y) - X
   {

diff  --git a/llvm/test/Transforms/InstCombine/exact.ll b/llvm/test/Transforms/InstCombine/exact.ll
index f4eb86cadc70a..ccad6f974ffb1 100644
--- a/llvm/test/Transforms/InstCombine/exact.ll
+++ b/llvm/test/Transforms/InstCombine/exact.ll
@@ -338,9 +338,8 @@ define <2 x i1> @sdiv_icmp6_vec(<2 x i64> %X) {
 
 define i8 @mul_of_udiv(i8 %x) {
 ; CHECK-LABEL: @mul_of_udiv(
-; CHECK-NEXT:    [[DIV:%.*]] = udiv exact i8 [[X:%.*]], 12
-; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[DIV]], 6
-; CHECK-NEXT:    ret i8 [[MUL]]
+; CHECK-NEXT:    [[MUL1:%.*]] = lshr exact i8 [[X:%.*]], 1
+; CHECK-NEXT:    ret i8 [[MUL1]]
 ;
   %div = udiv exact i8 %x, 12
   %mul = mul i8 %div, 6
@@ -349,8 +348,8 @@ define i8 @mul_of_udiv(i8 %x) {
 
 define i8 @mul_of_sdiv(i8 %x) {
 ; CHECK-LABEL: @mul_of_sdiv(
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv exact i8 [[X:%.*]], 12
-; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[DIV]], -6
+; CHECK-NEXT:    [[MUL_NEG:%.*]] = ashr exact i8 [[X:%.*]], 1
+; CHECK-NEXT:    [[MUL:%.*]] = sub nsw i8 0, [[MUL_NEG]]
 ; CHECK-NEXT:    ret i8 [[MUL]]
 ;
   %div = sdiv exact i8 %x, 12
@@ -360,8 +359,7 @@ define i8 @mul_of_sdiv(i8 %x) {
 
 define <2 x i8> @mul_of_sdiv_non_splat(<2 x i8> %x) {
 ; CHECK-LABEL: @mul_of_sdiv_non_splat(
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv exact <2 x i8> [[X:%.*]], <i8 6, i8 -12>
-; CHECK-NEXT:    [[MUL:%.*]] = mul <2 x i8> [[DIV]], <i8 6, i8 6>
+; CHECK-NEXT:    [[MUL:%.*]] = sdiv exact <2 x i8> [[X:%.*]], <i8 1, i8 -2>
 ; CHECK-NEXT:    ret <2 x i8> [[MUL]]
 ;
   %div = sdiv exact <2 x i8> %x, <i8 6, i8 -12>