[TableGen] Add support constraining a vector type in a pattern to have a specific element type and for constraining a vector type to have the same number of elements as another vector type. This is useful for AVX512 mask operations so we relate the mask type to the type of the other arguments.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@231356 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/Target/TargetSelectionDAG.td

@@ -68,6 +68,18 @@ class SDTCisSubVecOfVec<int ThisOp, int OtherOp>
   int OtherOpNum = OtherOp;
 }
 
+// SDTCVecEltisVT - The specified operand is vector type with element type
+// of VT.
+class SDTCVecEltisVT<int OpNum, ValueType vt> : SDTypeConstraint<OpNum> {
+  ValueType VT = vt;
+}
+
+// SDTCisSameNumEltsAs - The two specified operands have identical number
+// of elements.
+class SDTCisSameNumEltsAs<int OpNum, int OtherOp> : SDTypeConstraint<OpNum> {
+  int OtherOperandNum = OtherOp;
+}
+
 //===----------------------------------------------------------------------===//
 // Selection DAG Type Profile definitions.
 //

utils/TableGen/CodeGenDAGPatterns.cpp

@@ -482,6 +482,34 @@ bool EEVT::TypeSet::EnforceSmallerThan(EEVT::TypeSet &Other, TreePattern &TP) {
   return MadeChange;
 }
 
+/// EnforceVectorEltTypeIs - 'this' is now constrainted to be a vector type
+/// whose element is specified by VTOperand.
+bool EEVT::TypeSet::EnforceVectorEltTypeIs(MVT::SimpleValueType VT,
+                                           TreePattern &TP) {
+  bool MadeChange = false;
+
+  MadeChange |= EnforceVector(TP);
+
+  TypeSet InputSet(*this);
+
+  // Filter out all the types which don't have the right element type.
+  for (unsigned i = 0; i != TypeVec.size(); ++i) {
+    assert(isVector(TypeVec[i]) && "EnforceVector didn't work");
+    if (MVT(TypeVec[i]).getVectorElementType().SimpleTy != VT) {
+      TypeVec.erase(TypeVec.begin()+i--);
+      MadeChange = true;
+    }
+  }
+
+  if (TypeVec.empty()) {  // FIXME: Really want an SMLoc here!
+    TP.error("Type inference contradiction found, forcing '" +
+             InputSet.getName() + "' to have a vector element");
+    return false;
+  }
+
+  return MadeChange;
+}
+
 /// EnforceVectorEltTypeIs - 'this' is now constrainted to be a vector type
 /// whose element is specified by VTOperand.
 bool EEVT::TypeSet::EnforceVectorEltTypeIs(EEVT::TypeSet &VTOperand,
@@ -609,6 +637,64 @@ bool EEVT::TypeSet::EnforceVectorSubVectorTypeIs(EEVT::TypeSet &VTOperand,
   return MadeChange;
 }
 
+/// EnforceVectorSameNumElts - 'this' is now constrainted to
+/// be a vector with same num elements as VTOperand.
+bool EEVT::TypeSet::EnforceVectorSameNumElts(EEVT::TypeSet &VTOperand,
+                                             TreePattern &TP) {
+  if (TP.hasError())
+    return false;
+
+  // "This" must be a vector and "VTOperand" must be a vector.
+  bool MadeChange = false;
+  MadeChange |= EnforceVector(TP);
+  MadeChange |= VTOperand.EnforceVector(TP);
+
+  // If we know one of the vector types, it forces the other type to agree.
+  if (isConcrete()) {
+    MVT IVT = getConcrete();
+    unsigned NumElems = IVT.getVectorNumElements();
+
+    // Only keep types that have same elements as VTOperand.
+    TypeSet InputSet(VTOperand);
+
+    for (unsigned i = 0; i != VTOperand.TypeVec.size(); ++i) {
+      assert(isVector(VTOperand.TypeVec[i]) && "EnforceVector didn't work");
+      if (MVT(VTOperand.TypeVec[i]).getVectorNumElements() != NumElems) {
+        VTOperand.TypeVec.erase(VTOperand.TypeVec.begin()+i--);
+        MadeChange = true;
+      }
+    }
+    if (VTOperand.TypeVec.empty()) {  // FIXME: Really want an SMLoc here!
+      TP.error("Type inference contradiction found, forcing '" +
+               InputSet.getName() + "' to have same number elements as '" +
+               getName() + "'");
+      return false;
+    }
+  } else if (VTOperand.isConcrete()) {
+    MVT IVT = VTOperand.getConcrete();
+    unsigned NumElems = IVT.getVectorNumElements();
+
+    // Only keep types that have same elements as 'this'.
+    TypeSet InputSet(*this);
+
+    for (unsigned i = 0; i != TypeVec.size(); ++i) {
+      assert(isVector(TypeVec[i]) && "EnforceVector didn't work");
+      if (MVT(TypeVec[i]).getVectorNumElements() != NumElems) {
+        TypeVec.erase(TypeVec.begin()+i--);
+        MadeChange = true;
+      }
+    }
+    if (TypeVec.empty()) {  // FIXME: Really want an SMLoc here!
+      TP.error("Type inference contradiction found, forcing '" +
+               InputSet.getName() + "' to have same number elements than '" +
+               VTOperand.getName() + "'");
+      return false;
+    }
+  }
+
+  return MadeChange;
+}
+
 //===----------------------------------------------------------------------===//
 // Helpers for working with extended types.
 
@@ -839,6 +925,19 @@ SDTypeConstraint::SDTypeConstraint(Record *R) {
     ConstraintType = SDTCisSubVecOfVec;
     x.SDTCisSubVecOfVec_Info.OtherOperandNum =
       R->getValueAsInt("OtherOpNum");
+  } else if (R->isSubClassOf("SDTCVecEltisVT")) {
+    ConstraintType = SDTCVecEltisVT;
+    x.SDTCVecEltisVT_Info.VT = getValueType(R->getValueAsDef("VT"));
+    if (MVT(x.SDTCVecEltisVT_Info.VT).isVector())
+      PrintFatalError(R->getLoc(), "Cannot use vector type as SDTCVecEltisVT");
+    if (!MVT(x.SDTCVecEltisVT_Info.VT).isInteger() &&
+        !MVT(x.SDTCVecEltisVT_Info.VT).isFloatingPoint())
+      PrintFatalError(R->getLoc(), "Must use integer or floating point type "
+                                   "as SDTCVecEltisVT");
+  } else if (R->isSubClassOf("SDTCisSameNumEltsAs")) {
+    ConstraintType = SDTCisSameNumEltsAs;
+    x.SDTCisSameNumEltsAs_Info.OtherOperandNum =
+      R->getValueAsInt("OtherOperandNum");
   } else {
     errs() << "Unrecognized SDTypeConstraint '" << R->getName() << "'!\n";
     exit(1);
@@ -956,6 +1055,18 @@ bool SDTypeConstraint::ApplyTypeConstraint(TreePatternNode *N,
     return BigVecOperand->getExtType(VResNo).
       EnforceVectorSubVectorTypeIs(NodeToApply->getExtType(ResNo), TP);
   }
+  case SDTCVecEltisVT: {
+    return NodeToApply->getExtType(ResNo).
+      EnforceVectorEltTypeIs(x.SDTCVecEltisVT_Info.VT, TP);
+  }
+  case SDTCisSameNumEltsAs: {
+    unsigned OResNo = 0;
+    TreePatternNode *OtherNode =
+      getOperandNum(x.SDTCisSameNumEltsAs_Info.OtherOperandNum,
+                    N, NodeInfo, OResNo);
+    return OtherNode->getExtType(OResNo).
+      EnforceVectorSameNumElts(NodeToApply->getExtType(ResNo), TP);
+  }
   }
   llvm_unreachable("Invalid ConstraintType!");
 }

utils/TableGen/CodeGenDAGPatterns.h

@@ -136,10 +136,18 @@ namespace EEVT {
     /// whose element is VT.
     bool EnforceVectorEltTypeIs(EEVT::TypeSet &VT, TreePattern &TP);
 
+    /// EnforceVectorEltTypeIs - 'this' is now constrainted to be a vector type
+    /// whose element is VT.
+    bool EnforceVectorEltTypeIs(MVT::SimpleValueType VT, TreePattern &TP);
+
     /// EnforceVectorSubVectorTypeIs - 'this' is now constrainted to
     /// be a vector type VT.
     bool EnforceVectorSubVectorTypeIs(EEVT::TypeSet &VT, TreePattern &TP);
 
+    /// EnforceVectorSameNumElts - 'this' is now constrainted to
+    /// be a vector with same num elements as VT.
+    bool EnforceVectorSameNumElts(EEVT::TypeSet &VT, TreePattern &TP);
+
     bool operator!=(const TypeSet &RHS) const { return TypeVec != RHS.TypeVec; }
     bool operator==(const TypeSet &RHS) const { return TypeVec == RHS.TypeVec; }
 
@@ -165,7 +173,7 @@ struct SDTypeConstraint {
   enum {
     SDTCisVT, SDTCisPtrTy, SDTCisInt, SDTCisFP, SDTCisVec, SDTCisSameAs,
     SDTCisVTSmallerThanOp, SDTCisOpSmallerThanOp, SDTCisEltOfVec,
-    SDTCisSubVecOfVec
+    SDTCisSubVecOfVec, SDTCVecEltisVT, SDTCisSameNumEltsAs
   } ConstraintType;
 
   union {   // The discriminated union.
@@ -187,6 +195,12 @@ struct SDTypeConstraint {
     struct {
       unsigned OtherOperandNum;
     } SDTCisSubVecOfVec_Info;
+    struct {
+      MVT::SimpleValueType VT;
+    } SDTCVecEltisVT_Info;
+    struct {
+      unsigned OtherOperandNum;
+    } SDTCisSameNumEltsAs_Info;
   } x;
 
   /// ApplyTypeConstraint - Given a node in a pattern, apply this type