[TargetLowering] Add support for non-uniform vectors to BuildExactSDIV

This patch refactors the existing BuildExactSDIV implementation to support non-uniform constant vector denominators.

Differential Revision: https://reviews.llvm.org/D50392

llvm-svn: 339756

1 files changed, 24 insertions, 12 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index edb1b3e2fe7..5e3d284dcff 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -3438,32 +3438,44 @@ static SDValue BuildExactSDIV(const TargetLowering &TLI, SDNode *N,
   SDValue Op0 = N->getOperand(0);
   SDValue Op1 = N->getOperand(1);
   EVT VT = N->getValueType(0);
+  EVT SVT = VT.getScalarType();
   EVT ShVT = TLI.getShiftAmountTy(VT, DAG.getDataLayout());
+  EVT ShSVT = ShVT.getScalarType();
 
-  auto BuildSDIVPattern = [](APInt Divisor, unsigned &Shift, APInt &Factor) {
-    bool UseSRA = false;
-    Shift = Divisor.countTrailingZeros();
+  bool UseSRA = false;
+  SmallVector<SDValue, 16> Shifts, Factors;
+
+  auto BuildSDIVPattern = [&](ConstantSDNode *C) {
+    if (C->isNullValue())
+      return false;
+    APInt Divisor = C->getAPIntValue();
+    unsigned Shift = Divisor.countTrailingZeros();
     if (Shift) {
       Divisor.ashrInPlace(Shift);
       UseSRA = true;
     }
     // Calculate the multiplicative inverse, using Newton's method.
     APInt t;
-    Factor = Divisor;
+    APInt Factor = Divisor;
     while ((t = Divisor * Factor) != 1)
       Factor *= APInt(Divisor.getBitWidth(), 2) - t;
-    return UseSRA;
+    Shifts.push_back(DAG.getConstant(Shift, dl, ShSVT));
+    Factors.push_back(DAG.getConstant(Factor, dl, SVT));
+    return true;
   };
 
-  ConstantSDNode *C = isConstOrConstSplat(Op1);
-  if (!C || C->isNullValue())
+  // Collect all magic values from the build vector.
+  if (!ISD::matchUnaryPredicate(Op1, BuildSDIVPattern))
     return SDValue();
 
-  APInt FactorVal;
-  unsigned ShiftVal;
-  bool UseSRA = BuildSDIVPattern(C->getAPIntValue(), ShiftVal, FactorVal);
-  SDValue Shift = DAG.getConstant(ShiftVal, dl, ShVT);
-  SDValue Factor = DAG.getConstant(FactorVal, dl, VT);
+  SDValue Shift, Factor;
+  if (VT.isVector()) {
+    Shift = DAG.getBuildVector(ShVT, dl, Shifts);
+    Factor = DAG.getBuildVector(VT, dl, Factors);
+  } else {
+    Shift = Shifts[0];
+    Factor = Factors[0];
+  }
 
   SDValue Res = Op0;