[InstCombine] Add fold for `(rem (mul/shl X, Y), (mul/shl X, Z))` -> `(mul X, (rem Y, Z))`

[goldstein.w.n]

Summary

Forked from D142901 to handle the mul/shl when Y and Z are
non-constant.

(srem (mul/shl X, Y), (mul/shl X, Z))

[If `Y`/`Z` are one use]
    -> `(mul/shl nuw nsw X, (srem Y, Z))`
        - https://alive2.llvm.org/ce/z/ccTFj2
        - https://alive2.llvm.org/ce/z/i_UQ5A
    -> `(mul/shl nsw X, (srem Y, Z))`
        - https://alive2.llvm.org/ce/z/mQKc63
        - https://alive2.llvm.org/ce/z/uERkKH

(urem (mul/shl X, Y), (mul/shl X, Z))

[If `Y`/`Z` are one use]
    -> `(mul/shl nuw nsw X, (srem Y, Z))`
        - https://alive2.llvm.org/ce/z/mNnQqJ
        - https://alive2.llvm.org/ce/z/Bj_DR-
        - https://alive2.llvm.org/ce/z/X6ZEtQ
    -> `(mul/shl nuw X, (srem Y, Z))`
        - https://alive2.llvm.org/ce/z/SJYtUV

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

ping.

[goldstein.w.n]

Rebase + bug fix

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase - no change

[goldstein.w.n]

Rebase

[goldstein.w.n]

ping @sdesmalen any update any when you will have time review this (no rush).

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

static Instruction *simplifyIRemMulShl(BinaryOperator &I,
   bool IsSigned = I.getOpcode() == Instruction::SRem;
 
   bool NSW0 = BO0->hasNoSignedWrap();
   bool NSW1 = BO1->hasNoSignedWrap();
   bool NUW0 = BO0->hasNoUnsignedWrap();
   bool NUW1 = BO1->hasNoUnsignedWrap();
 
   // Try to handle constant cases first.
-  if (!ConstY || !ConstZ)
-    return nullptr;
-
+  if (ConstY && ConstZ) {
     APInt APIntY = ConstY->getValue();
     APInt APIntZ = ConstZ->getValue();
 
     // Just treat the shifts as mul, we may end up returning a mul by power
     // of 2 but that will be cleaned up later.
     if (BO0->getOpcode() == Instruction::Shl)
       APIntY = APInt(APIntY.getBitWidth(), 1) << APIntY;
     if (BO1->getOpcode() == Instruction::Shl)
       APIntZ = APInt(APIntZ.getBitWidth(), 1) << APIntZ;
 
     APInt RemYZ = IsSigned ? APIntY.srem(APIntZ) : APIntY.urem(APIntZ);
 
     // (rem (mul nuw/nsw X, Y), (mul X, Z))
     //      if (rem Y, Z) == 0
     //          -> 0
     if (RemYZ.isZero() && (IsSigned ? NSW0 : NUW0))
       return IC.replaceInstUsesWith(I, ConstantInt::getNullValue(Ty));
 
     // (rem (mul X, Y), (mul nuw/nsw X, Z))
     //      if (rem Y, Z) == Y
     //          -> (mul nuw/nsw X, Y)
     if (RemYZ == APIntY && (IsSigned ? NSW1 : NUW1)) {
       // We are returning Op0 essentially but we can also add no wrap flags.
       BinaryOperator *BO =
           BinaryOperator::CreateMul(X, ConstantInt::get(Ty, APIntY));
       // We can add nsw/nuw if remainder op is signed/unsigned, also we
       // can copy any overflow flags from Op0.
       if (IsSigned || NSW0)
         BO->setHasNoSignedWrap();
       if (!IsSigned || NUW0)
         BO->setHasNoUnsignedWrap();
       return BO;
     }
 
     // (rem (mul nuw/nsw X, Y), (mul {nsw} X, Z))
     //      if Y >= Z
     //          -> (mul {nuw} nsw X, (rem Y, Z))
     // NB: (rem Y, Z) is a constant.
     if (APIntY.uge(APIntZ) && (IsSigned ? (NSW0 && NSW1) : NUW0)) {
       BinaryOperator *BO =
           BinaryOperator::CreateMul(X, ConstantInt::get(Ty, RemYZ));
       BO->setHasNoSignedWrap();
       if (!IsSigned || NUW0)
         BO->setHasNoUnsignedWrap();
       return BO;
     }
+  }
+
+  // Check if desirable to do generic replacement.
+  // NB: It may be beneficial to do this if we have X << Z even if there are
+  // multiple uses of Op0/Op1 as it will eliminate the urem (urem of a power
+  // of 2 is converted to add/and) and urem is pretty expensive (maybe more
+  // sense in DAGCombiner).
+  if ((ConstY && ConstZ) ||
+      (Op0->hasOneUse() && Op1->hasOneUse() &&
+       (IsSigned ? (BO0->getOpcode() != Instruction::Shl &&
+                    BO1->getOpcode() != Instruction::Shl)
+                 : (BO0->getOpcode() != Instruction::Shl ||
+                    BO1->getOpcode() == Instruction::Shl)))) {
+    // (rem (mul nuw/nsw X, Y), (mul nuw {nsw} X, Z)
+    //        -> (mul nuw/nsw X, (rem Y, Z))
+    if (IsSigned ? (NSW0 && NSW1 && NUW1) : (NUW0 && NUW1)) {
+      // Convert the shifts to multiplies, cleaned up elsewhere.
+      if (BO0->getOpcode() == Instruction::Shl)
+        Y = IC.Builder.CreateShl(ConstantInt::get(Ty, 1), Y);
+      if (BO1->getOpcode() == Instruction::Shl)
+        Z = IC.Builder.CreateShl(ConstantInt::get(Ty, 1), Z);
+      BinaryOperator *BO =
+          BinaryOperator::CreateMul(X, IsSigned ? IC.Builder.CreateSRem(Y, Z)
+                                                : IC.Builder.CreateURem(Y, Z));
+
+      if (IsSigned || NSW0 || NSW1)
+        BO->setHasNoSignedWrap();
+      if (!IsSigned || (NUW0 && NUW1))
+        BO->setHasNoUnsignedWrap();
+      return BO;
+    }
+  }
 
   return nullptr;
 }
 
 /// This function implements the transforms common to both integer remainder
 /// instructions (urem and srem). It is called by the visitors to those integer
 /// remainder instructions.
 /// Common integer remainder transforms

llvm/test/Transforms/InstCombine/urem-mul.ll

;
   %BO0 = mul nsw i8 %X, 21
   %BO1 = mul nsw nuw i8 %X, 6
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_Y_Z_is_mul_X_RemYZ(i8 %X, i8 %Y, i8 %Z) {
 ; CHECK-LABEL: @urem_Y_Z_is_mul_X_RemYZ(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nuw i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = mul nuw i8 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = urem i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw i8 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw i8 %X, %Y
   %BO1 = mul nuw i8 %Z, %X
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_CX_Y_Z_is_mul_X_RemYZ(i8 %X, i8 %Y, i8 %Z) {
 ; CHECK-LABEL: @urem_CX_Y_Z_is_mul_X_RemYZ(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nuw i8 [[Y:%.*]], 10
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw i8 10, [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw i8 -1, [[Z:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 [[NOTMASK]], -1
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw i8 [[TMP2]], 10
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw i8 10, %Y
   %BO1 = shl nuw i8 10, %Z
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_Y_Z_is_mul_X_RemYZ_with_nsw_out1(i8 %X, i8 %Y, i8 %Z) {
 ; CHECK-LABEL: @urem_Y_Z_is_mul_X_RemYZ_with_nsw_out1(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nuw nsw i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw i8 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw i8 -1, [[Z:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 [[NOTMASK]], -1
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i8 [[TMP2]], [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw nsw i8 %X, %Y
   %BO1 = shl nuw i8 %X, %Z
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define <2 x i8> @urem_Y_Z_is_mul_X_RemYZ_with_nsw_out2(<2 x i8> %X, <2 x i8> %Y, <2 x i8> %Z) {
 ; CHECK-LABEL: @urem_Y_Z_is_mul_X_RemYZ_with_nsw_out2(
-; CHECK-NEXT:    [[BO0:%.*]] = shl nuw <2 x i8> [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw nsw <2 x i8> [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = urem <2 x i8> [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw <2 x i8> <i8 1, i8 1>, [[Y:%.*]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw <2 x i8> <i8 -1, i8 -1>, [[Z:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = xor <2 x i8> [[NOTMASK]], <i8 -1, i8 -1>
+; CHECK-NEXT:    [[TMP3:%.*]] = and <2 x i8> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw <2 x i8> [[TMP3]], [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %BO0 = shl nuw <2 x i8> %Y, %X
   %BO1 = shl nuw nsw <2 x i8> %Z, %X
   %r = urem <2 x i8> %BO0, %BO1
   ret <2 x i8> %r
 }
 
;
   %BO0 = shl nuw i8 %X, 4
   %BO1 = mul nsw nuw i8 %X, 5
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_Y_Z_is_mul_X_RemYZ(i8 %X, i8 %Y, i8 %Z) {
 ; CHECK-LABEL: @srem_Y_Z_is_mul_X_RemYZ(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nsw i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = mul nuw nsw i8 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = srem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = srem i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul nsw i8 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nsw i8 %Y, %X
   %BO1 = mul nsw nuw i8 %X, %Z
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_Y_Z_is_mul_X_RemYZ_with_nuw_out(i8 %X, i8 %Y, i8 %Z) {
 ; CHECK-LABEL: @srem_Y_Z_is_mul_X_RemYZ_with_nuw_out(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nuw nsw i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = mul nuw nsw i8 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = srem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = srem i8 [[Y:%.*]], [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i8 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nsw nuw i8 %Y, %X
   %BO1 = mul nsw nuw i8 %Z, %X
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }