[InstCombine] Add transforms for `(rem (shl Y, X), (shl Z, X))`

[goldstein.w.n]

Summary

This is just filling in a missing case from D144225.

We treat (shl Y, X) and (shl Z, X) as (mul Z, 1 << X) and `(mul
Y, 1 << X)` then reuse the same transformations that already exist.

[goldstein.w.n]

Rebase (with cleaner pattern match)

[goldstein.w.n]

ping.

[goldstein.w.n]

ping2.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1762–1767

Can you add this new case to the description here?

1771

I wonder if you can do something similar to MatchShiftOrMul, so that we use a similar mechanism

if (MatchShiftOrMul(Op0, X, Y) &&
    MatchShiftOrMul(Op1, X, Z, /*MatchSpecific=*/true))
  ; // Do nothing
else if (MatchShiftConstByVal(Op0, Y, X) &&
         MatchShiftConstByVal(Op1, Z, X, /*MatchSpecific=*/true))
  ShiftX = true;
else
  return nullptr;

Just note that this requires changing the logic in MatchShiftOrMul which currently checks if V is nullptr to use m_Specific only when requested through an operand to the function. This way, you don't need to 'reset' V in between the two calls. I'm actually thinking that this explicit behaviour may be more desirable anyway, because then there is no implicit dependence between successive calls of this function.

1828

Can you add a comment describing the need for this function?

1828

nit: is it worth renaming this to CreateMulOrShift ?

1829

nit: rename to ShiftByX ?

This might be overkill because there are only two options, but for readability purposes I wondered if it was worth adding an enum which describes the two options, e.g. enum RemMatchKind { NoMatch = 0, MulValByConst = 1, ConstShiftedByVal = 2 }; or something.

1843–1845

A number of these comments are now out of date. You could add a comment which describes that:

C << X <=> C * (1 << X)

so that the same logic here still applies when you conceptually substitute (1 << X) by X.
It just requires that the return value must be implemented differently (which is why there is the GetBinOpOut function)

1847

Can you pass in the APInt instead, and do the ConstantInt::get(...) in GetBinOptOut?

[goldstein.w.n]

Rebase

[sdesmalen]

Thanks for the changes, LGTM!

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

nit: unnecessary curly braces (same above)

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

IIRC, style is if any of the if need a curly, then all should have.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

In this case none of the blocks needs the curly braces if you change the first block to ; // pass ?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

In this case none of the blocks needs the curly braces if you change the first block to ; // pass ?

Sure, will do. I'll update when I rebase later. Any chance you will also be able to review D143417?

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

Thanks! Are you happy to merge the three accepted patches?

I've currently got quite a few things on my plate so I'm not sure how much time I'll have to review D143417 (especially given its complexity), but I'll try to have a look at it next week.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1808

Sure, I'll merge inbetween reviews. No rush and thank you for the reviews thus far :)

[sdesmalen]

Ping, gentle reminder to land these patches.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Value *Y1 =
         Builder.CreateFAddFMF(Y, ConstantFP::get(I.getType(), -1.0), &I);
     Value *Pow = Builder.CreateBinaryIntrinsic(Intrinsic::pow, Op1, Y1, &I);
     return replaceInstUsesWith(I, Pow);
   }
 
   return nullptr;
 }
 
-// Variety of transform for (urem/srem (mul/shl X, Y), (mul/shl X, Z))
+// Variety of transform for:
+//  (urem/srem (mul X, Y), (mul X, Z))
+//  (urem/srem (shl X, Y), (shl X, Z))
+//  (urem/srem (shl Y, X), (shl Z, X))
+// NB: The shift cases are really just extensions of the mul case. We treat
+// shift as Val * (1 << Amt).
 static Instruction *simplifyIRemMulShl(BinaryOperator &I,
                                        InstCombinerImpl &IC) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1), *X = nullptr;
   APInt Y, Z;
+  bool ShiftByX = false;
 
   // If V is not nullptr, it will be matched using m_Specific.
-  auto MatchShiftOrMul = [](Value *Op, Value *&V, APInt &C) -> bool {
+  auto MatchShiftOrMulXC = [](Value *Op, Value *&V, APInt &C) -> bool {
     const APInt *Tmp = nullptr;
     if ((!V && match(Op, m_Mul(m_Value(V), m_APInt(Tmp)))) ||
         (V && match(Op, m_Mul(m_Specific(V), m_APInt(Tmp)))))
       C = *Tmp;
     else if ((!V && match(Op, m_Shl(m_Value(V), m_APInt(Tmp)))) ||
              (V && match(Op, m_Shl(m_Specific(V), m_APInt(Tmp)))))
       C = APInt(Tmp->getBitWidth(), 1) << *Tmp;
-    return Tmp != nullptr;
+    if (Tmp != nullptr)
+      return true;
+
+    // Reset `V` so we don't start with specific value on next match attempt.
+    V = nullptr;
+    return false;
   };
 
-  if (!MatchShiftOrMul(Op0, X, Y) || !MatchShiftOrMul(Op1, X, Z))
+  auto MatchShiftCX = [](Value *Op, APInt &C, Value *&V) -> bool {
+    const APInt *Tmp = nullptr;
+    if ((!V && match(Op, m_Shl(m_APInt(Tmp), m_Value(V)))) ||
+        (V && match(Op, m_Shl(m_APInt(Tmp), m_Specific(V))))) {
+      C = *Tmp;
+      return true;
+    }
+
+    // Reset `V` so we don't start with specific value on next match attempt.
+    V = nullptr;
+    return false;
+  };
+
+  if (MatchShiftOrMulXC(Op0, X, Y) && MatchShiftOrMulXC(Op1, X, Z)) {
+    // pass
+  } else if (MatchShiftCX(Op0, Y, X) && MatchShiftCX(Op1, Z, X)) {
+    ShiftByX = true;
+  } else {
     return nullptr;
+  }
 
   bool IsSRem = I.getOpcode() == Instruction::SRem;
 
   OverflowingBinaryOperator *BO0 = cast<OverflowingBinaryOperator>(Op0);
   // TODO: We may be able to deduce more about nsw/nuw of BO0/BO1 based on Y >=
   // Z or Z >= Y.
   bool BO0HasNSW = BO0->hasNoSignedWrap();
   bool BO0HasNUW = BO0->hasNoUnsignedWrap();
   bool BO0NoWrap = IsSRem ? BO0HasNSW : BO0HasNUW;
 
   APInt RemYZ = IsSRem ? Y.srem(Z) : Y.urem(Z);
   // (rem (mul nuw/nsw X, Y), (mul X, Z))
   //      if (rem Y, Z) == 0
   //          -> 0
   if (RemYZ.isZero() && BO0NoWrap)
     return IC.replaceInstUsesWith(I, ConstantInt::getNullValue(I.getType()));
 
+  // Helper function to emit either (RemSimplificationC << X) or
+  // (RemSimplificationC * X) depending on whether we matched Op0/Op1 as
+  // (shl V, X) or (mul V, X) respectively.
+  auto CreateMulOrShift =
+      [&](const APInt &RemSimplificationC) -> BinaryOperator * {
+    Value *RemSimplification =
+        ConstantInt::get(I.getType(), RemSimplificationC);
+    return ShiftByX ? BinaryOperator::CreateShl(RemSimplification, X)
+                    : BinaryOperator::CreateMul(X, RemSimplification);
+  };
+
   OverflowingBinaryOperator *BO1 = cast<OverflowingBinaryOperator>(Op1);
   bool BO1HasNSW = BO1->hasNoSignedWrap();
   bool BO1HasNUW = BO1->hasNoUnsignedWrap();
   bool BO1NoWrap = IsSRem ? BO1HasNSW : BO1HasNUW;
   // (rem (mul X, Y), (mul nuw/nsw X, Z))
   //      if (rem Y, Z) == Y
   //          -> (mul nuw/nsw X, Y)
   if (RemYZ == Y && BO1NoWrap) {
-    BinaryOperator *BO =
-        BinaryOperator::CreateMul(X, ConstantInt::get(I.getType(), Y));
+    BinaryOperator *BO = CreateMulOrShift(Y);
     // Copy any overflow flags from Op0.
     BO->setHasNoSignedWrap(IsSRem || BO0HasNSW);
     BO->setHasNoUnsignedWrap(!IsSRem || BO0HasNUW);
     return BO;
   }
 
   // (rem (mul nuw/nsw X, Y), (mul {nsw} X, Z))
   //      if Y >= Z
   //          -> (mul {nuw} nsw X, (rem Y, Z))
   if (Y.uge(Z) && (IsSRem ? (BO0HasNSW && BO1HasNSW) : BO0HasNUW)) {
-    BinaryOperator *BO =
-        BinaryOperator::CreateMul(X, ConstantInt::get(I.getType(), RemYZ));
+    BinaryOperator *BO = CreateMulOrShift(RemYZ);
     BO->setHasNoSignedWrap();
     BO->setHasNoUnsignedWrap(BO0HasNUW);
     return BO;
   }
 
   return nullptr;
 }
 

llvm/test/Transforms/InstCombine/rem-mul-shl.ll

;
   %BO0 = mul nuw i8 %X, 15
   %BO1 = mul i8 %X, 5
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_XY_XZ_with_CY_rem_CZ_eq_0_with_shl(i8 %X) {
 ; CHECK-LABEL: @urem_XY_XZ_with_CY_rem_CZ_eq_0_with_shl(
-; CHECK-NEXT:    [[BO0:%.*]] = shl nuw i8 15, [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = shl i8 5, [[X]]
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 0
 ;
   %BO0 = shl nuw i8 15, %X
   %BO1 = shl i8 5, %X
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_XY_XZ_with_CY_rem_CZ_eq_0_fail_missing_flag(i8 %X) {
;
   %BO0 = mul i8 %X, 3
   %BO1 = mul nuw i8 %X, 12
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @urem_XY_XZ_with_CY_lt_CZ_with_shl(i8 %X) {
 ; CHECK-LABEL: @urem_XY_XZ_with_CY_lt_CZ_with_shl(
-; CHECK-NEXT:    [[BO0:%.*]] = shl i8 3, [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw i8 12, [[X]]
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i8 3, [[X:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = shl i8 3, %X
   %BO1 = shl nuw i8 12, %X
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
;
   %BO0 = mul nuw i8 %X, 5
   %BO1 = mul nsw i8 %X, 15
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define <2 x i8> @srem_XY_XZ_with_CY_lt_CZ_with_nuw_out_with_shl(<2 x i8> %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_lt_CZ_with_nuw_out_with_shl(
-; CHECK-NEXT:    [[BO0:%.*]] = shl nuw <2 x i8> <i8 3, i8 3>, [[X:%.*]]
-; CHECK-NEXT:    [[BO1:%.*]] = shl nsw <2 x i8> <i8 15, i8 15>, [[X]]
-; CHECK-NEXT:    [[R:%.*]] = srem <2 x i8> [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw nsw <2 x i8> <i8 3, i8 3>, [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %BO0 = shl nuw <2 x i8> <i8 3, i8 3>, %X
   %BO1 = shl nsw <2 x i8> <i8 15, i8 15>, %X
   %r = srem <2 x i8> %BO0, %BO1
   ret <2 x i8> %r
 }