[InstCombine] Add constant combines for `(urem/srem (shl X, Y), (shl X, Z))`

[goldstein.w.n]

Summary

Forked from D142901 to deduce more nsw/nuw flag for the output
shl.

We can handle the following cases + some nsw/nuw flags:

The rationale for doing this all in InstCombine rather than handling
the constant shl cases in InstSimplify is we often create a new
instruction because we are able to deduce more nsw/nuw flags than
the original instruction had.

[goldstein.w.n]

Use match instead of handwritten logic

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1689

It might be a bit easier to follow, if you explicitly do the scaling while doing the matching, i.e.

APInt Y, Z;
const APInt *MatchY = nullptr, *MatchZ = nullptr;

// Match and normalise shift-amounts to multiplications
if (match(Op0, m_c_Mul(m_Value(X), m_APInt(MatchY)))) {
  Y = *MatchY;
  if (match(Op1, m_Shl(m_Specific(X), m_APInt(MatchZ))))
    //  rem(mul(x, y), shl(x, z))
    Z = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchZ);
  else if (match(Op1, m_c_Mul(m_Specific(X), m_APInt(MatchZ))))
    //  rem(mul(x, y), mul(x, z))
    Z = *MatchZ;
} else if (match(Op0, m_Shl(m_Value(X), m_APInt(MatchY)))) {
  Y = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchY);
  if (match(Op1, m_Shl(m_Specific(X), m_APInt(MatchZ))))
    //  rem(shl(x, y), shl(x, z))
    Z = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchZ);
  else if (match(Op1, m_c_Mul(m_Specific(X), m_APInt(MatchZ))))
    //  rem(shl(x, y), mul(x, z))
    Z = *MatchZ;
}

if (!MatchY || !MatchZ)
  return nullptr;

1704–1706

I can't see this case being tested anywhere. I'd suggest singling this case out from the other logic and handling it separately. Maybe you can move that to another patch (or make it part of D143417, which tries to handle a similar case). That makes this patch a bit simpler and avoids the need for GetBinOpOut.

[goldstein.w.n]

Split of (shl Z, X), (shl Y, X) case. Remove the Shift{Y|Z|X} flags

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1689

It might be a bit easier to follow, if you explicitly do the scaling while doing the matching, i.e.

Inlined the ShiftY/ShiftZ

Prefer keeping the 4 explicit cases (removed the ShiftX case and moved to next patch). Think its clearer to have the cases each explicitly laid out, rather than having nest if/else statements. LMK if thats okay, will change if you feel strongly.

APInt Y, Z;
const APInt *MatchY = nullptr, *MatchZ = nullptr;

// Match and normalise shift-amounts to multiplications
if (match(Op0, m_c_Mul(m_Value(X), m_APInt(MatchY)))) {
  Y = *MatchY;
  if (match(Op1, m_Shl(m_Specific(X), m_APInt(MatchZ))))
    //  rem(mul(x, y), shl(x, z))
    Z = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchZ);
  else if (match(Op1, m_c_Mul(m_Specific(X), m_APInt(MatchZ))))
    //  rem(mul(x, y), mul(x, z))
    Z = *MatchZ;
} else if (match(Op0, m_Shl(m_Value(X), m_APInt(MatchY)))) {
  Y = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchY);
  if (match(Op1, m_Shl(m_Specific(X), m_APInt(MatchZ))))
    //  rem(shl(x, y), shl(x, z))
    Z = APInt(X->getType()->getScalarSizeInBits(), 1).shl(*MatchZ);
  else if (match(Op1, m_c_Mul(m_Specific(X), m_APInt(MatchZ))))
    //  rem(shl(x, y), mul(x, z))
    Z = *MatchZ;
}

if (!MatchY || !MatchZ)
  return nullptr;

1704–1706

I can't see this case being tested anywhere. I'd suggest singling this case out from the other logic and handling it separately. Maybe you can move that to another patch (or make it part of D143417, which tries to handle a similar case). That makes this patch a bit simpler and avoids the need for GetBinOpOut.

There are some tests for it, though they are all non-constant cases (so it starts to matter in the next patch).

I split it to a new patch (see test=D147107, impl=D147108) and I added tests for the constant version of case.

[goldstein.w.n]

@sdesmalen ping

[goldstein.w.n]

ping @sdesmalen

[sdesmalen]

LGTM with comment about redundant condition addressed.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1689

This is merely a suggestion, I'll leave it to you whether to adopt.

From looking at the new way you've structured the code, it occurred to me that it can also be written as this:

// If V is not nullptr, it will be matched using m_Specific.
auto MatchShiftOrMul = [](Value *Op, Value *&V, APInt &C) -> bool {
  const APInt *Tmp = nullptr;
  if ((!V && match(Op, m_c_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;
};

APInt Y, Z;
Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1), *X = nullptr;
if (!MatchShiftOrMul(Op0, X, Y) || !MatchShiftOrMul(Op1, X, Z))
  return nullptr;

Which avoids having to spell out all the permutations. It also avoids the need for the AdjustedY and AdjustedZ variables.

1697–1699

You can remove this condition, because InstCombine will already have canonicalised the constant to the RHS.

[goldstein.w.n]

Use cleaner method for matching

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1689

Done, likewise for D147108. But couldn't find a clean way to do it for
D143417. In D143417 its not longer only APInt matches so need
branching logic for either V or Y/Z (Tmp) being nullptr.

[goldstein.w.n]

Rebase

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Instruction *InstCombinerImpl::visitFDiv(BinaryOperator &I) {
 
   return nullptr;
 }
 
 // Variety of transform for (urem/srem (mul/shl X, Y), (mul/shl X, Z))
 static Instruction *simplifyIRemMulShl(BinaryOperator &I,
                                        InstCombinerImpl &IC) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1), *X, *Y, *Z;
-  if (!(match(Op0, m_Mul(m_Value(X), m_Value(Y))) &&
-        match(Op1, m_c_Mul(m_Specific(X), m_Value(Z)))) &&
-      !(match(Op0, m_Mul(m_Value(Y), m_Value(X))) &&
-        match(Op1, m_c_Mul(m_Specific(X), m_Value(Z)))))
+  bool ShiftX = false, ShiftY = false, ShiftZ = false;
+  if ((match(Op0, m_Mul(m_Value(X), m_Value(Y))) &&
+       match(Op1, m_c_Mul(m_Specific(X), m_Value(Z)))) ||
+      (match(Op0, m_Mul(m_Value(Y), m_Value(X))) &&
+       match(Op1, m_c_Mul(m_Specific(X), m_Value(Z))))) {
+    // Pass
+  } else if (match(Op0, m_Shl(m_Value(X), m_Value(Y))) &&
+             match(Op1, m_c_Mul(m_Specific(X), m_Value(Z)))) {
+    ShiftY = true;
+  } else if (match(Op1, m_Shl(m_Value(X), m_Value(Z))) &&
+             match(Op0, m_c_Mul(m_Specific(X), m_Value(Y)))) {
+    ShiftZ = true;
+  } else if (match(Op0, m_Shl(m_Value(X), m_Value(Y))) &&
+             match(Op1, m_Shl(m_Specific(X), m_Value(Z)))) {
+    ShiftZ = true;
+    ShiftY = true;
+  } else if (match(Op0, m_Shl(m_Value(Y), m_Value(X))) &&
+             match(Op1, m_Shl(m_Value(Z), m_Specific(X)))) {
+    ShiftX = true;
+  } else {
     return nullptr;
+  }
+
+  BinaryOperator *BO0 = cast<BinaryOperator>(Op0);
+  BinaryOperator *BO1 = cast<BinaryOperator>(Op1);
 
   auto GetOperandAsConstantInt = [](Value *Op) -> ConstantInt * {
     if (Op->getType()->isVectorTy())
       if (auto *COp = dyn_cast<Constant>(Op)) {
         auto *CSplat = COp->getSplatValue();
         return CSplat ? dyn_cast<ConstantInt>(CSplat) : nullptr;
       }
     return dyn_cast<ConstantInt>(Op);
   };
 
   ConstantInt *ConstY = GetOperandAsConstantInt(Y);
   ConstantInt *ConstZ = GetOperandAsConstantInt(Z);
 
   bool IsSRem = I.getOpcode() == Instruction::SRem;
 
   if (!ConstY || !ConstZ)
     return nullptr;
 
   APInt APIntY = ConstY->getValue();
   APInt APIntZ = ConstZ->getValue();
 
-  BinaryOperator *BO0 = cast<BinaryOperator>(Op0);
-  BinaryOperator *BO1 = cast<BinaryOperator>(Op1);
+  // 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 (ShiftY)
+    APIntY = APInt(APIntY.getBitWidth(), 1) << APIntY;
+  if (ShiftZ)
+    APIntZ = APInt(APIntZ.getBitWidth(), 1) << APIntZ;
 
   bool BO0HasNSW = BO0->hasNoSignedWrap();
   bool BO0HasNUW = BO0->hasNoUnsignedWrap();
 
   APInt RemYZ = IsSRem ? APIntY.srem(APIntZ) : APIntY.urem(APIntZ);
   // (rem (mul nuw/nsw X, Y), (mul X, Z))
   //      if (rem Y, Z) == 0
   //          -> 0
   if (RemYZ.isZero() && (IsSRem ? BO0HasNSW : BO0HasNUW))
     return IC.replaceInstUsesWith(I, ConstantInt::getNullValue(I.getType()));
 
+  auto GetBinOpOut = [&](Value *RemSimplification) -> BinaryOperator * {
+    return ShiftX ? BinaryOperator::CreateShl(RemSimplification, X)
+                  : BinaryOperator::CreateMul(X, RemSimplification);
+  };
+
   bool BO1HasNSW = BO1->hasNoSignedWrap();
   bool BO1HasNUW = BO1->hasNoUnsignedWrap();
 
   // (rem (mul X, Y), (mul nuw/nsw X, Z))
   //      if (rem Y, Z) == Y
   //          -> (mul nuw/nsw X, Y)
   if (RemYZ == APIntY && (IsSRem ? BO1HasNSW : BO1HasNUW)) {
-    BinaryOperator *BO =
-        BinaryOperator::CreateMul(X, ConstantInt::get(I.getType(), APIntY));
+    BinaryOperator *BO = GetBinOpOut(ConstantInt::get(I.getType(), APIntY));
     // Copy any overflow flags from Op0.
     if (IsSRem || BO0HasNSW)
       BO->setHasNoSignedWrap();
     if (!IsSRem || BO0HasNUW)
       BO->setHasNoUnsignedWrap();
     return BO;
   }
 
   // (rem (mul nuw/nsw X, Y), (mul {nsw} X, Z))
   //      if Y >= Z
   //          -> (mul {nuw} nsw X, (rem Y, Z))
   if (APIntY.uge(APIntZ) && (IsSRem ? (BO0HasNSW && BO1HasNSW) : BO0HasNUW)) {
-    BinaryOperator *BO =
-        BinaryOperator::CreateMul(X, ConstantInt::get(I.getType(), RemYZ));
+    BinaryOperator *BO = GetBinOpOut(ConstantInt::get(I.getType(), RemYZ));
     BO->setHasNoSignedWrap();
     if (BO0HasNUW)
       BO->setHasNoUnsignedWrap();
     return BO;
   }
 
   return nullptr;
 }

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

;
   %BO0 = mul i8 %X, 3
   %BO1 = mul nuw i8 %X, 12
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define <2 x i8> @urem_XY_XZ_with_CY_lt_CZ_with_nsw_out(<2 x i8> %X) {
 ; CHECK-LABEL: @urem_XY_XZ_with_CY_lt_CZ_with_nsw_out(
-; CHECK-NEXT:    [[BO0:%.*]] = shl nsw <2 x i8> [[X:%.*]], <i8 2, i8 2>
-; CHECK-NEXT:    [[BO1:%.*]] = mul nuw <2 x i8> [[X]], <i8 12, i8 12>
-; CHECK-NEXT:    [[R:%.*]] = urem <2 x i8> [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw nsw <2 x i8> [[X:%.*]], <i8 2, i8 2>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %BO0 = shl nsw <2 x i8> %X, <i8 2, i8 2>
   %BO1 = mul nuw <2 x i8> %X, <i8 12, i8 12>
   %r = urem <2 x i8> %BO0, %BO1
   ret <2 x i8> %r
 }
 
 define i8 @urem_XY_XZ_with_CY_lt_CZ_no_nsw_out(i8 %X) {
 ; CHECK-LABEL: @urem_XY_XZ_with_CY_lt_CZ_no_nsw_out(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nuw i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw nsw i8 [[X]], 3
-; CHECK-NEXT:    [[R:%.*]] = urem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw i8 [[X:%.*]], 3
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw i8 %X, 3
   %BO1 = shl nsw nuw i8 %X, 3
   %r = urem i8 %BO0, %BO1
   ret i8 %r
 }
 
;
   %BO0 = mul nuw i8 %X, 9
   %BO1 = mul nsw nuw i8 %X, 3
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define <2 x i8> @srem_XY_XZ_with_CY_lt_CZ(<2 x i8> %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_lt_CZ(
-; CHECK-NEXT:    [[BO0:%.*]] = shl <2 x i8> [[X:%.*]], <i8 3, i8 3>
-; CHECK-NEXT:    [[BO1:%.*]] = mul nsw <2 x i8> [[X]], <i8 15, i8 15>
-; CHECK-NEXT:    [[R:%.*]] = srem <2 x i8> [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nsw <2 x i8> [[X:%.*]], <i8 3, i8 3>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %BO0 = shl <2 x i8> %X, <i8 3, i8 3>
   %BO1 = mul nsw <2 x i8> %X, <i8 15, i8 15>
   %r = srem <2 x i8> %BO0, %BO1
   ret <2 x i8> %r
 }
 
 define i8 @srem_XY_XZ_with_CY_lt_CZ_with_nuw_out(i8 %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_lt_CZ_with_nuw_out(
 ; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i8 [[X:%.*]], 5
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw i8 %X, 5
   %BO1 = mul nsw i8 %X, 15
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_XY_XZ_with_CY_lt_CZ_no_nsw_out(i8 %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_lt_CZ_no_nsw_out(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nsw i8 [[X:%.*]], 5
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw nsw i8 [[X]], 4
-; CHECK-NEXT:    [[R:%.*]] = srem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = mul nsw i8 [[X:%.*]], 5
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nsw i8 %X, 5
   %BO1 = shl nsw nuw i8 %X, 4
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_XY_XZ_with_CY_lt_CZ_fail_missing_flag(i8 %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_lt_CZ_fail_missing_flag(
 ; CHECK-NEXT:    [[BO0:%.*]] = mul nuw nsw i8 [[X:%.*]], 5
 ; CHECK-NEXT:    [[BO1:%.*]] = shl nuw i8 [[X]], 4
 ; CHECK-NEXT:    [[R:%.*]] = srem i8 [[BO0]], [[BO1]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nuw nsw i8 %X, 5
   %BO1 = shl nuw i8 %X, 4
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_XY_XZ_with_CY_gt_CZ(i8 %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_gt_CZ(
-; CHECK-NEXT:    [[BO0:%.*]] = shl nsw i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[BO1:%.*]] = mul nsw i8 [[X]], 6
-; CHECK-NEXT:    [[R:%.*]] = srem i8 [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nsw i8 [[X:%.*]], 1
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = shl nsw i8 %X, 3
   %BO1 = mul nsw i8 %X, 6
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define i8 @srem_XY_XZ_with_CY_gt_CZ_with_nuw_out(i8 %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_gt_CZ_with_nuw_out(
 ; CHECK-NEXT:    [[R:%.*]] = shl nuw nsw i8 [[X:%.*]], 2
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %BO0 = mul nsw nuw i8 %X, 10
   %BO1 = mul nsw i8 %X, 6
   %r = srem i8 %BO0, %BO1
   ret i8 %r
 }
 
 define <2 x i8> @srem_XY_XZ_with_CY_gt_CZ_no_nuw_out(<2 x i8> %X) {
 ; CHECK-LABEL: @srem_XY_XZ_with_CY_gt_CZ_no_nuw_out(
-; CHECK-NEXT:    [[BO0:%.*]] = mul nsw <2 x i8> [[X:%.*]], <i8 10, i8 10>
-; CHECK-NEXT:    [[BO1:%.*]] = shl nuw nsw <2 x i8> [[X]], <i8 3, i8 3>
-; CHECK-NEXT:    [[R:%.*]] = srem <2 x i8> [[BO0]], [[BO1]]
+; CHECK-NEXT:    [[R:%.*]] = shl nsw <2 x i8> [[X:%.*]], <i8 1, i8 1>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %BO0 = mul nsw <2 x i8> %X, <i8 10, i8 10>
   %BO1 = shl nsw nuw <2 x i8> %X, <i8 3, i8 3>
   %r = srem <2 x i8> %BO0, %BO1
   ret <2 x i8> %r
 }