From ffc0217bc7f5fab0150bf537774859006d30492a Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Tue, 16 Jul 2019 11:54:54 -0700
Subject: [PATCH] Add a TypeIsPred.

Mostly one would use the type specification directly on the operand, but for
cases where the type of the operand depends on other operand types, `TypeIs`
attribute can be used to construct verification methods.

PiperOrigin-RevId: 258411758
---
 mlir/include/mlir/IR/OpBase.td       | 7 +++++++
 mlir/test/lib/TestDialect/TestOps.td | 9 +++++++++
 mlir/test/mlir-tblgen/types.mlir     | 8 ++++++++
 3 files changed, 24 insertions(+)

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 5ccd2a3..ef22681 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1172,6 +1172,13 @@ class TCopVTEtIs<int idx, Type type> : And<[
 
 // Predicate to verify that a named argument or result's element type matches a
 // given type.
+class TypeIsPred<string name, Type type> :
+   SubstLeaves<"$_self", "$" # name # ".getType()", type.predicate>;
+class TypeIs<string name, Type type> : PredOpTrait<
+  "'" # name # "' is " # type.description, TypeIsPred<name, type>>;
+
+// Predicate to verify that a named argument or result's element type matches a
+// given type.
 class ElementTypeIsPred<string name, Type type> : And<[
    SubstLeaves<"$_self", "$" # name # ".getType()", IsShapedTypePred>,
    SubstLeaves<"$_self", "getElementTypeOrSelf($" # name # ")",
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index e4ecc17..2640d41 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -187,6 +187,15 @@ def OperandOneAndResultHaveSameType :
   let results = (outs AnyTensor:$res);
 }
 
+def IfFirstOperandIsNoneThenSoIsSecond :
+    TEST_Op<"if_first_operand_is_none_then_so_is_second", [PredOpTrait<
+    "has either both none type operands or first is not none",
+     Or<[
+        And<[TypeIsPred<"x", NoneType>, TypeIsPred<"y", NoneType>]>,
+        Neg<TypeIsPred<"x", NoneType>>]>>]> {
+  let arguments = (ins AnyType:$x, AnyType:$y);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index 73fc7fb..4a94468 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -174,3 +174,11 @@ func @same_types_shape_mismatch(%arg0: tensor<1x2xi32>, %arg1: tensor<2x1xi32>)
   %0 = "test.operand_one_and_result_have_same_type"(%arg0, %arg1) : (tensor<1x2xi32>, tensor<2x1xi32>) -> tensor<2x1xi32>
   return
 }
+
+// -----
+
+func @does_not_have_i32(%arg0: tensor<1x2xi32>, %arg1: none) {
+  // expected-error@+1 {{either both none type operands or first is not none}}
+  "test.if_first_operand_is_none_then_so_is_second"(%arg1, %arg0) : (none, tensor<1x2xi32>) -> ()
+  return
+}
-- 
2.7.4