Support for type constraints across operand and results

--

PiperOrigin-RevId: 240905555

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7a441ef..7dfc2a3 100644 (file)
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -698,6 +698,20 @@ class TCopVTEtIsSameAs<int i, int j> : AllOf<[
     // TODO: This could be made into C++ function instead.
     CPred<"{0}.getOperand(" # i # ")->getType().cast<VectorOrTensorType>().getElementType() == {0}.getOperand(" # j # ")->getType().cast<VectorOrTensorType>().getElementType()">]>;
 
+// Predicate to verify that the i'th result and the j'th operand have the same
+// elemental type.
+// Type Constraint result`i`'s Vector or Tensor Element type is Same As
+// Type Constraint Operand `j`'s Vector or Tensor Element type.
+class TCresVTEtIsSameAsOp<int i, int j> : AllOf<[
+    CPred<"{0}.getNumResults() > " # i>,
+    CPred<"{0}.getNumOperands() > " # j>,
+    SubstLeaves<"{0}",  "{0}.getResult(" # i # ")->getType()",
+      IsVectorOrTensorTypePred>,
+    SubstLeaves<"{0}",  "{0}.getOperand(" # j # ")->getType()",
+      IsVectorOrTensorTypePred>,
+    // TODO: This could be made into C++ function instead.
+    CPred<"{0}.getResult(" # i # ")->getType().cast<VectorOrTensorType>().getElementType() == {0}.getOperand(" # j # ")->getType().cast<VectorOrTensorType>().getElementType()">]>;
+
 //===----------------------------------------------------------------------===//
 // Pattern definitions
 //===----------------------------------------------------------------------===//

diff --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
index 9dabf79..d777699 100644 (file)
--- a/mlir/test/mlir-tblgen/predicate.td
+++ b/mlir/test/mlir-tblgen/predicate.td
@@ -16,7 +16,11 @@ def TautologyElementTest: PredOpTrait<
   "first operand is a vector or tensor with the same elemental type as itself",
   TCopVTEtIsSameAs<0, 0>>;
 
-def Identity : Op<"identity", [TautologyElementTest]> {
+def OperandAndResultElementTypeTest : PredOpTrait<
+   "first operand is a vector or tensor with the same elemental type as first result",
+   TCresVTEtIsSameAsOp<0, 0>>;
+
+def Identity : Op<"identity", [TautologyElementTest, OperandAndResultElementTypeTest]> {
   let arguments = (ins
       I32OrF32Tensor:$x);
   let results = (outs
@@ -39,6 +43,9 @@ def IdentityI32 : Op<"identity_i32", [PredOpTrait<
 // Verify tautology constraint.
 // CHECK: if (!((((*this->getOperation()).getNumOperands() > std::max(0,0))) && (((*this->getOperation()).getOperand(0)->getType().isa<VectorOrTensorType>())) && (((*this->getOperation()).getOperand(0)->getType().isa<VectorOrTensorType>())) && (((*this->getOperation()).getOperand(0)->getType().cast<VectorOrTensorType>().getElementType() == (*this->getOperation()).getOperand(0)->getType().cast<VectorOrTensorType>().getElementType()))))
 // CHECK-NEXT: return emitOpError("failed to verify that first operand is a vector or tensor with the same elemental type as itself");
+// Verify OperandAndResultElementTypeTest constraint
+// CHECK: if (!((((*this->getOperation()).getNumResults() > 0)) && (((*this->getOperation()).getNumOperands() > 0)) && (((*this->getOperation()).getResult(0)->getType().isa<VectorOrTensorType>())) && (((*this->getOperation()).getOperand(0)->getType().isa<VectorOrTensorType>())) && (((*this->getOperation()).getResult(0)->getType().cast<VectorOrTensorType>().getElementType() == (*this->getOperation()).getOperand(0)->getType().cast<VectorOrTensorType>().getElementType()))))
+// CHECK-NEXT: return emitOpError("failed to verify that first operand is a vector or tensor with the same elemental type as first result");
 
 // CHECK-LABEL: IdentityI32::verify
 // CHECK: if (!((((*this->getOperation()).getNumOperands() > 0)) && (((*this->getOperation()).getOperand(0)->getType().isa<VectorOrTensorType>())) && (((*this->getOperation()).getOperand(0)->getType().cast<VectorOrTensorType>().getElementType().isInteger(32)))))