1 files changed, 31 insertions, 41 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 5e3d284dcff..e12ec54a2c4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -3581,19 +3581,27 @@ SDValue TargetLowering::BuildUDIV(SDNode *N, SelectionDAG &DAG,
                                   SmallVectorImpl<SDNode *> &Created) const {
   SDLoc dl(N);
   EVT VT = N->getValueType(0);
+  EVT SVT = VT.getScalarType();
   EVT ShVT = getShiftAmountTy(VT, DAG.getDataLayout());
+  EVT ShSVT = ShVT.getScalarType();
+  unsigned EltBits = VT.getScalarSizeInBits();
 
   // Check to see if we can do this.
   // FIXME: We should be more aggressive here.
   if (!isTypeLegal(VT))
     return SDValue();
 
-  auto BuildUDIVPattern = [](const APInt &Divisor, unsigned &PreShift,
-                             APInt &Magic, unsigned &PostShift) {
+  bool UseNPQ = false;
+  SmallVector<SDValue, 16> PreShifts, PostShifts, MagicFactors, NPQFactors;
+
+  auto BuildUDIVPattern = [&](ConstantSDNode *C) {
+    if (C->isNullValue())
+      return false;
     // FIXME: We should use a narrower constant when the upper
     // bits are known to be zero.
+    APInt Divisor = C->getAPIntValue();
     APInt::mu magics = Divisor.magicu();
-    PreShift = PostShift = 0;
+    unsigned PreShift = 0, PostShift = 0;
 
     // If the divisor is even, we can avoid using the expensive fixup by
     // shifting the divided value upfront.
@@ -3604,65 +3612,47 @@ SDValue TargetLowering::BuildUDIV(SDNode *N, SelectionDAG &DAG,
       assert(magics.a == 0 && "Should use cheap fixup now");
     }
 
-    Magic = magics.m;
+    APInt Magic = magics.m;
 
+    unsigned SelNPQ;
     if (magics.a == 0 || Divisor.isOneValue()) {
       assert(magics.s < Divisor.getBitWidth() &&
              "We shouldn't generate an undefined shift!");
       PostShift = magics.s;
-      return false;
+      SelNPQ = false;
     } else {
       PostShift = magics.s - 1;
-      return true;
+      SelNPQ = true;
     }
+
+    PreShifts.push_back(DAG.getConstant(PreShift, dl, ShSVT));
+    MagicFactors.push_back(DAG.getConstant(Magic, dl, SVT));
+    NPQFactors.push_back(
+        DAG.getConstant(SelNPQ ? APInt::getOneBitSet(EltBits, EltBits - 1)
+                               : APInt::getNullValue(EltBits),
+                        dl, SVT));
+    PostShifts.push_back(DAG.getConstant(PostShift, dl, ShSVT));
+    UseNPQ |= SelNPQ;
+    return true;
   };
 
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
 
   // Collect the shifts/magic values from each element.
-  bool UseNPQ = false;
+  if (!ISD::matchUnaryPredicate(N1, BuildUDIVPattern))
+    return SDValue();
+
   SDValue PreShift, PostShift, MagicFactor, NPQFactor;
   if (VT.isVector()) {
-    EVT SVT = VT.getScalarType();
-    EVT ShSVT = ShVT.getScalarType();
-    unsigned EltBits = VT.getScalarSizeInBits();
-    unsigned NumElts = VT.getVectorNumElements();
-    SmallVector<SDValue, 16> PreShifts, PostShifts, MagicFactors, NPQFactors;
-    if (ISD::BUILD_VECTOR != N1.getOpcode())
-      return SDValue();
-    for (unsigned i = 0; i != NumElts; ++i) {
-      auto *C = dyn_cast<ConstantSDNode>(N1.getOperand(i));
-      if (!C || C->isNullValue() || C->getAPIntValue().getBitWidth() != EltBits)
-        return SDValue();
-      APInt MagicVal;
-      unsigned PreShiftVal, PostShiftVal;
-      bool SelNPQ = BuildUDIVPattern(C->getAPIntValue(), PreShiftVal, MagicVal,
-                                     PostShiftVal);
-      PreShifts.push_back(DAG.getConstant(PreShiftVal, dl, ShSVT));
-      MagicFactors.push_back(DAG.getConstant(MagicVal, dl, SVT));
-      NPQFactors.push_back(
-          DAG.getConstant(SelNPQ ? APInt::getOneBitSet(EltBits, EltBits - 1)
-                                 : APInt::getNullValue(EltBits),
-                          dl, SVT));
-      PostShifts.push_back(DAG.getConstant(PostShiftVal, dl, ShSVT));
-      UseNPQ |= SelNPQ;
-    }
     PreShift = DAG.getBuildVector(ShVT, dl, PreShifts);
     MagicFactor = DAG.getBuildVector(VT, dl, MagicFactors);
     NPQFactor = DAG.getBuildVector(VT, dl, NPQFactors);
     PostShift = DAG.getBuildVector(ShVT, dl, PostShifts);
   } else {
-    auto *C = dyn_cast<ConstantSDNode>(N1);
-    if (!C || C->isNullValue())
-      return SDValue();
-    APInt MagicVal;
-    unsigned PreShiftVal, PostShiftVal;
-    UseNPQ = BuildUDIVPattern(C->getAPIntValue(), PreShiftVal, MagicVal,
-                              PostShiftVal);
-    PreShift = DAG.getConstant(PreShiftVal, dl, ShVT);
-    MagicFactor = DAG.getConstant(MagicVal, dl, VT);
-    PostShift = DAG.getConstant(PostShiftVal, dl, ShVT);
+    PreShift = PreShifts[0];
+    MagicFactor = MagicFactors[0];
+    PostShift = PostShifts[0];
   }
 
   SDValue Q = N0;