[InstCombine] fix and enhance udiv/urem narrowing

[spatel]

Summary

There are 3 independent changes here, but I thought it would be easier to review one patch to handle them all. If it's better to review and/or commit separately, please let me know:

1. Account for multiple uses in the pattern matching: avoid the transform if it increases the instruction count.
2. Add a missing fold for the case where the numerator is the constant: http://rise4fun.com/Alive/E2p
3. Enable all folds for vector types.

There's still one more potential change - use "shouldChangeType()" to keep from transforming to an illegal integer type.

[craig.topper]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1103

This is allowing arbitrary constant expressions through. Is that we want?

[majnemer]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1103

I don't see an issue with allowing ConstantExpr thru.

[spatel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1103

The intent was of course to let arbitrary vector constants through, but it should work for constexpr too. I'll add a test to make sure.

[spatel]

Patch updated:
No code changes, but added a test with a constexpr to show that's ok.

[craig.topper]

LGTM

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (match(Op0, m_Sub(m_Value(X), m_Value(Z)))) { // (X - Z) / Y; Y = Op1
     if ((isSigned && match(Z, m_SRem(m_Specific(X), m_Specific(Op1)))) ||
         (!isSigned && match(Z, m_URem(m_Specific(X), m_Specific(Op1)))))
       return BinaryOperator::Create(I.getOpcode(), X, Op1);
   }
 
   return nullptr;
 }
 
-/// dyn_castZExtVal - Checks if V is a zext or constant that can
-/// be truncated to Ty without losing bits.
-static Value *dyn_castZExtVal(Value *V, Type *Ty) {
-  if (ZExtInst *Z = dyn_cast<ZExtInst>(V)) {
-    if (Z->getSrcTy() == Ty)
-      return Z->getOperand(0);
-  } else if (ConstantInt *C = dyn_cast<ConstantInt>(V)) {
-    if (C->getValue().getActiveBits() <= cast<IntegerType>(Ty)->getBitWidth())
-      return ConstantExpr::getTrunc(C, Ty);
-  }
-  return nullptr;
-}
-
 namespace {
 const unsigned MaxDepth = 6;
 typedef Instruction *(*FoldUDivOperandCb)(Value *Op0, Value *Op1,
                                           const BinaryOperator &I,
                                           InstCombiner &IC);
 
 /// \brief Used to maintain state for visitUDivOperand().
 struct UDivFoldAction {
if (size_t LHSIdx =
       if (visitUDivOperand(Op0, SI->getOperand(2), I, Actions, Depth)) {
         Actions.push_back(UDivFoldAction(nullptr, Op1, LHSIdx - 1));
         return Actions.size();
       }
 
   return 0;
 }
 
+/// If we have zero-extended operands of an unsigned div or rem, we may be able
+/// to narrow the operation (sink the zext below the math).
+static Instruction *narrowUDivURem(BinaryOperator &I,
+                                   InstCombiner::BuilderTy &Builder) {
+  Instruction::BinaryOps Opcode = I.getOpcode();
+  Value *N = I.getOperand(0);
+  Value *D = I.getOperand(1);
+  Type *Ty = I.getType();
+  Value *X, *Y;
+  if (match(N, m_ZExt(m_Value(X))) && match(D, m_ZExt(m_Value(Y))) &&
+      X->getType() == Y->getType() && (N->hasOneUse() || D->hasOneUse())) {
+    // udiv (zext X), (zext Y) --> zext (udiv X, Y)
+    // urem (zext X), (zext Y) --> zext (urem X, Y)
+    Value *NarrowOp = Builder.CreateBinOp(Opcode, X, Y);
+    return new ZExtInst(NarrowOp, Ty);
+  }
+
+  Constant *C;
+  if ((match(N, m_OneUse(m_ZExt(m_Value(X)))) && match(D, m_Constant(C))) ||
+      (match(D, m_OneUse(m_ZExt(m_Value(X)))) && match(N, m_Constant(C)))) {
+    // If the constant is the same in the smaller type, use the narrow version.
+    Constant *TruncC = ConstantExpr::getTrunc(C, X->getType());
+    if (ConstantExpr::getZExt(TruncC, Ty) != C)
+      return nullptr;
+
+    // udiv (zext X), C --> zext (udiv X, C')
+    // urem (zext X), C --> zext (urem X, C')
+    // udiv C, (zext X) --> zext (udiv C', X)
+    // urem C, (zext X) --> zext (urem C', X)
+    Value *NarrowOp = isa<Constant>(D) ? Builder.CreateBinOp(Opcode, X, TruncC)
+                                       : Builder.CreateBinOp(Opcode, TruncC, X);
+    return new ZExtInst(NarrowOp, Ty);
+  }
+
+  return nullptr;
+}
+
 Instruction *InstCombiner::visitUDiv(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
 
   if (Value *V = SimplifyVectorOp(I))
     return replaceInstUsesWith(I, V);
 
   if (Value *V = SimplifyUDivInst(Op0, Op1, SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
if (match(Op0, m_LShr(m_Value(X), m_APInt(C1))) &&
             X, ConstantInt::get(X->getType(), C2ShlC1));
         if (IsExact)
           BO->setIsExact();
         return BO;
       }
     }
   }
 
-  // (zext A) udiv (zext B) --> zext (A udiv B)
-  if (ZExtInst *ZOp0 = dyn_cast<ZExtInst>(Op0))
-    if (Value *ZOp1 = dyn_castZExtVal(Op1, ZOp0->getSrcTy()))
-      return new ZExtInst(
-          Builder.CreateUDiv(ZOp0->getOperand(0), ZOp1, "div", I.isExact()),
-          I.getType());
+  if (Instruction *NarrowDiv = narrowUDivURem(I, Builder))
+    return NarrowDiv;
 
   // (LHS udiv (select (select (...)))) -> (LHS >> (select (select (...))))
   SmallVector<UDivFoldAction, 6> UDivActions;
   if (visitUDivOperand(Op0, Op1, I, UDivActions))
     for (unsigned i = 0, e = UDivActions.size(); i != e; ++i) {
       FoldUDivOperandCb Action = UDivActions[i].FoldAction;
       Value *ActionOp1 = UDivActions[i].OperandToFold;
       Instruction *Inst;
if (Value *V = SimplifyVectorOp(I))
     return replaceInstUsesWith(I, V);
 
   if (Value *V = SimplifyURemInst(Op0, Op1, SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (Instruction *common = commonIRemTransforms(I))
     return common;
 
-  // (zext A) urem (zext B) --> zext (A urem B)
-  if (ZExtInst *ZOp0 = dyn_cast<ZExtInst>(Op0))
-    if (Value *ZOp1 = dyn_castZExtVal(Op1, ZOp0->getSrcTy()))
-      return new ZExtInst(Builder.CreateURem(ZOp0->getOperand(0), ZOp1),
-                          I.getType());
+  if (Instruction *NarrowRem = narrowUDivURem(I, Builder))
+    return NarrowRem;
 
   // X urem Y -> X and Y-1, where Y is a power of 2,
   if (isKnownToBeAPowerOfTwo(Op1, /*OrZero*/ true, 0, &I)) {
     Constant *N1 = Constant::getAllOnesValue(I.getType());
     Value *Add = Builder.CreateAdd(Op1, N1);
     return BinaryOperator::CreateAnd(Op0, Add);
   }
 

test/Transforms/InstCombine/udivrem-change-width.ll

;
   %udiv = udiv <2 x i32> %za, %zb
   ret <2 x i32> %udiv
 }
 
 define i32 @udiv_i32_multiuse(i8 %a, i8 %b) {
 ; CHECK-LABEL: @udiv_i32_multiuse(
 ; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
 ; CHECK-NEXT:    [[ZB:%.*]] = zext i8 %b to i32
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 %a, %b
-; CHECK-NEXT:    [[UDIV:%.*]] = zext i8 [[DIV]] to i32
+; CHECK-NEXT:    [[UDIV:%.*]] = udiv i32 [[ZA]], [[ZB]]
 ; CHECK-NEXT:    [[EXTRA_USES:%.*]] = add nuw nsw i32 [[ZA]], [[ZB]]
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[EXTRA_USES]], [[UDIV]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[UDIV]], [[EXTRA_USES]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %za = zext i8 %a to i32
   %zb = zext i8 %b to i32
   %udiv = udiv i32 %za, %zb
   %extra_uses = add i32 %za, %zb
   %r = mul i32 %udiv, %extra_uses
   ret i32 %r
;
   %urem = urem <2 x i32> %za, %zb
   ret <2 x i32> %urem
 }
 
 define i32 @urem_i32_multiuse(i8 %a, i8 %b) {
 ; CHECK-LABEL: @urem_i32_multiuse(
 ; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
 ; CHECK-NEXT:    [[ZB:%.*]] = zext i8 %b to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = urem i8 %a, %b
-; CHECK-NEXT:    [[UREM:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    [[UREM:%.*]] = urem i32 [[ZA]], [[ZB]]
 ; CHECK-NEXT:    [[EXTRA_USES:%.*]] = add nuw nsw i32 [[ZA]], [[ZB]]
-; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[EXTRA_USES]], [[UREM]]
+; CHECK-NEXT:    [[R:%.*]] = mul nuw nsw i32 [[UREM]], [[EXTRA_USES]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %za = zext i8 %a to i32
   %zb = zext i8 %b to i32
   %urem = urem i32 %za, %zb
   %extra_uses = add i32 %za, %zb
   %r = mul i32 %urem, %extra_uses
   ret i32 %r

 ;
   %za = zext i8 %a to i32
   %udiv = udiv i32 %za, 10
   ret i32 %udiv
 }
 
 define <2 x i32> @udiv_i32_c_vec(<2 x i8> %a) {
 ; CHECK-LABEL: @udiv_i32_c_vec(
-; CHECK-NEXT:    [[ZA:%.*]] = zext <2 x i8> %a to <2 x i32>
-; CHECK-NEXT:    [[UDIV:%.*]] = udiv <2 x i32> [[ZA]], <i32 10, i32 17>
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv <2 x i8> %a, <i8 10, i8 17>
+; CHECK-NEXT:    [[UDIV:%.*]] = zext <2 x i8> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[UDIV]]
 ;
   %za = zext <2 x i8> %a to <2 x i32>
   %udiv = udiv <2 x i32> %za, <i32 10, i32 17>
   ret <2 x i32> %udiv
 }
 
 define i32 @udiv_i32_c_multiuse(i8 %a) {
 ; CHECK-LABEL: @udiv_i32_c_multiuse(
 ; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 %a, 10
-; CHECK-NEXT:    [[UDIV:%.*]] = zext i8 [[DIV]] to i32
-; CHECK-NEXT:    [[EXTRA_USE:%.*]] = add nuw nsw i32 [[ZA]], [[UDIV]]
+; CHECK-NEXT:    [[UDIV:%.*]] = udiv i32 [[ZA]], 10
+; CHECK-NEXT:    [[EXTRA_USE:%.*]] = add nuw nsw i32 [[UDIV]], [[ZA]]
 ; CHECK-NEXT:    ret i32 [[EXTRA_USE]]
 ;
   %za = zext i8 %a to i32
   %udiv = udiv i32 %za, 10
   %extra_use = add i32 %za, %udiv
   ret i32 %extra_use
 }
 

 ;
   %za = zext i8 %a to i32
   %urem = urem i32 %za, 10
   ret i32 %urem
 }
 
 define <2 x i32> @urem_i32_c_vec(<2 x i8> %a) {
 ; CHECK-LABEL: @urem_i32_c_vec(
-; CHECK-NEXT:    [[ZA:%.*]] = zext <2 x i8> %a to <2 x i32>
-; CHECK-NEXT:    [[UREM:%.*]] = urem <2 x i32> [[ZA]], <i32 10, i32 17>
+; CHECK-NEXT:    [[TMP1:%.*]] = urem <2 x i8> %a, <i8 10, i8 17>
+; CHECK-NEXT:    [[UREM:%.*]] = zext <2 x i8> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[UREM]]
 ;
   %za = zext <2 x i8> %a to <2 x i32>
   %urem = urem <2 x i32> %za, <i32 10, i32 17>
   ret <2 x i32> %urem
 }
 
 define i32 @urem_i32_c_multiuse(i8 %a) {
 ; CHECK-LABEL: @urem_i32_c_multiuse(
 ; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = urem i8 %a, 10
-; CHECK-NEXT:    [[UREM:%.*]] = zext i8 [[TMP1]] to i32
-; CHECK-NEXT:    [[EXTRA_USE:%.*]] = add nuw nsw i32 [[ZA]], [[UREM]]
+; CHECK-NEXT:    [[UREM:%.*]] = urem i32 [[ZA]], 10
+; CHECK-NEXT:    [[EXTRA_USE:%.*]] = add nuw nsw i32 [[UREM]], [[ZA]]
 ; CHECK-NEXT:    ret i32 [[EXTRA_USE]]
 ;
   %za = zext i8 %a to i32
   %urem = urem i32 %za, 10
   %extra_use = add i32 %za, %urem
   ret i32 %extra_use
 }
 
 define i32 @urem_illegal_type_c(i9 %a) {
 ; CHECK-LABEL: @urem_illegal_type_c(
 ; CHECK-NEXT:    [[TMP1:%.*]] = urem i9 %a, 10
 ; CHECK-NEXT:    [[UREM:%.*]] = zext i9 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[UREM]]
 ;
   %za = zext i9 %a to i32
   %urem = urem i32 %za, 10
   ret i32 %urem
 }
 
 define i32 @udiv_c_i32(i8 %a) {
 ; CHECK-LABEL: @udiv_c_i32(
-; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
-; CHECK-NEXT:    [[UDIV:%.*]] = udiv i32 10, [[ZA]]
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv i8 10, %a
+; CHECK-NEXT:    [[UDIV:%.*]] = zext i8 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[UDIV]]
 ;
   %za = zext i8 %a to i32
   %udiv = udiv i32 10, %za
   ret i32 %udiv
 }
 
 define i32 @urem_c_i32(i8 %a) {
 ; CHECK-LABEL: @urem_c_i32(
-; CHECK-NEXT:    [[ZA:%.*]] = zext i8 %a to i32
-; CHECK-NEXT:    [[UREM:%.*]] = urem i32 10, [[ZA]]
+; CHECK-NEXT:    [[TMP1:%.*]] = urem i8 10, %a
+; CHECK-NEXT:    [[UREM:%.*]] = zext i8 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[UREM]]
 ;
   %za = zext i8 %a to i32
   %urem = urem i32 10, %za
   ret i32 %urem
 }