Fix getNode to allow a vector for the shift amount for shifts of vectors.

Fix the shift amount when unrolling a vector shift into scalar shifts.
Fix problem in getShuffleScalarElt where it assumes that the input of
a bit convert must be a vector.

llvm-svn: 60740

2 files changed, 33 insertions, 9 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
index d61d6ca6fdf..445d32ff1b8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
@@ -297,6 +297,9 @@ private:
 
   SDValue ExpandEXTRACT_SUBVECTOR(SDValue Op);
   SDValue ExpandEXTRACT_VECTOR_ELT(SDValue Op);
+
+  // Returns the legalized (truncated or extended) shift amount.
+  SDValue LegalizeShiftAmount(SDValue ShiftAmt);
 };
 }
 
@@ -786,8 +789,19 @@ SDValue SelectionDAGLegalize::UnrollVectorOp(SDValue Op) {
         Operands[j] = Operand;
       }
     }
-    Scalars.push_back(DAG.getNode(Op.getOpcode(), EltVT,
-                                  &Operands[0], Operands.size()));
+
+    switch (Op.getOpcode()) {
+    default:
+      Scalars.push_back(DAG.getNode(Op.getOpcode(), EltVT,
+                                    &Operands[0], Operands.size()));
+      break;
+    case ISD::SHL:
+    case ISD::SRA:
+    case ISD::SRL:
+      Scalars.push_back(DAG.getNode(Op.getOpcode(), EltVT, Operands[0],
+                                    LegalizeShiftAmount(Operands[1])));
+      break;
+    }
   }
 
   return DAG.getNode(ISD::BUILD_VECTOR, VT, &Scalars[0], Scalars.size());
@@ -850,6 +864,17 @@ PerformInsertVectorEltInMemory(SDValue Vec, SDValue Val, SDValue Idx) {
                      PseudoSourceValue::getFixedStack(SPFI), 0);
 }
 
+SDValue SelectionDAGLegalize::LegalizeShiftAmount(SDValue ShiftAmt) {
+  if (TLI.getShiftAmountTy().bitsLT(ShiftAmt.getValueType()))
+    return DAG.getNode(ISD::TRUNCATE, TLI.getShiftAmountTy(), ShiftAmt);
+
+  if (TLI.getShiftAmountTy().bitsGT(ShiftAmt.getValueType()))
+    return DAG.getNode(ISD::ANY_EXTEND, TLI.getShiftAmountTy(), ShiftAmt);
+
+  return ShiftAmt;
+}
+
+
 /// LegalizeOp - We know that the specified value has a legal type, and
 /// that its operands are legal.  Now ensure that the operation itself
 /// is legal, recursively ensuring that the operands' operations remain
@@ -3094,11 +3119,8 @@ SDValue SelectionDAGLegalize::LegalizeOp(SDValue Op) {
          Node->getOpcode() == ISD::SRL ||
          Node->getOpcode() == ISD::SRA) &&
         !Node->getValueType(0).isVector()) {
-      if (TLI.getShiftAmountTy().bitsLT(Tmp2.getValueType()))
-        Tmp2 = DAG.getNode(ISD::TRUNCATE, TLI.getShiftAmountTy(), Tmp2);
-      else if (TLI.getShiftAmountTy().bitsGT(Tmp2.getValueType()))
-        Tmp2 = DAG.getNode(ISD::ANY_EXTEND, TLI.getShiftAmountTy(), Tmp2);
-    }
+      Tmp2 = LegalizeShiftAmount(Tmp2);
+     }
 
     Result = DAG.UpdateNodeOperands(Result, Tmp1, Tmp2);
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 6fc2a67e296..34f0ccac47d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -2066,7 +2066,8 @@ SDValue SelectionDAG::getShuffleScalarElt(const SDNode *N, unsigned i) {
 
   if (V.getOpcode() == ISD::BIT_CONVERT) {
     V = V.getOperand(0);
-    if (V.getValueType().getVectorNumElements() != NumElems)
+    MVT VVT = V.getValueType();
+    if (!VVT.isVector() || VVT.getVectorNumElements() != NumElems)
       return SDValue();
   }
   if (V.getOpcode() == ISD::SCALAR_TO_VECTOR)
@@ -2418,7 +2419,8 @@ SDValue SelectionDAG::getNode(unsigned Opcode, MVT VT,
            "Shift operators return type must be the same as their first arg");
     assert(VT.isInteger() && N2.getValueType().isInteger() &&
            "Shifts only work on integers");
-    assert(N2.getValueType() == TLI.getShiftAmountTy() &&
+    assert((N2.getValueType() == TLI.getShiftAmountTy() ||
+            (N2.getValueType().isVector() && N2.getValueType().isInteger())) &&
            "Wrong type for shift amount");
 
     // Always fold shifts of i1 values so the code generator doesn't need to