[Mlir-commits] [mlir] e81bf67 - [MLIR] Modify HasParent trait to allow one of several op's as a parent

[Mehdi Amini]

Author: Rahul Joshi
Date: 2020-06-16T04:50:56Z
New Revision: e81bf67e8cf91ecfc8727fb03b9af20063f74d71

URL: https://github.com/llvm/llvm-project/commit/e81bf67e8cf91ecfc8727fb03b9af20063f74d71
DIFF: https://github.com/llvm/llvm-project/commit/e81bf67e8cf91ecfc8727fb03b9af20063f74d71.diff

LOG: [MLIR] Modify HasParent trait to allow one of several op's as a parent

- Modify HasParent trait to allow one of several op's as a parent -
- Expose this trait in the ODS framework using the ParentOneOf<> trait.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/include/mlir/IR/OpDefinition.h
    mlir/test/IR/traits.mlir
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6d05b78cbf9a..b1a28572e003 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1684,6 +1684,9 @@ class SingleBlockImplicitTerminator<string op>
 class HasParent<string op>
     : ParamNativeOpTrait<"HasParent", op>;
 
+class ParentOneOf<list<string> ops>
+    : ParamNativeOpTrait<"HasParent", StrJoin<ops>.result>;
+
 // Op result type is derived from the first attribute. If the attribute is an
 // subclass of `TypeAttrBase`, its value is used, otherwise, the type of the
 // attribute content is used.

diff  --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index da4912d3c1f5..3fc68d09f840 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1145,16 +1145,22 @@ template <typename TerminatorOpType> struct SingleBlockImplicitTerminator {
   };
 };
 
-/// This class provides a verifier for ops that are expecting a specific parent.
-template <typename ParentOpType> struct HasParent {
+/// This class provides a verifier for ops that are expecting their parent
+/// to be one of the given parent ops
+template <typename... ParentOpTypes>
+struct HasParent {
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType, Impl> {
   public:
     static LogicalResult verifyTrait(Operation *op) {
-      if (isa<ParentOpType>(op->getParentOp()))
+      if (llvm::isa<ParentOpTypes...>(op->getParentOp()))
         return success();
-      return op->emitOpError() << "expects parent op '"
-                               << ParentOpType::getOperationName() << "'";
+
+      return op->emitOpError()
+             << "expects parent op "
+             << (sizeof...(ParentOpTypes) != 1 ? "to be one of '" : "'")
+             << llvm::makeArrayRef({ParentOpTypes::getOperationName()...})
+             << "'";
     }
   };
 };

diff  --git a/mlir/test/IR/traits.mlir b/mlir/test/IR/traits.mlir
index a076d8f10308..a08a22ce1adc 100644
--- a/mlir/test/IR/traits.mlir
+++ b/mlir/test/IR/traits.mlir
@@ -173,6 +173,39 @@ func @failedHasParent_wrong_parent() {
   }) : () -> ()
 }
 
+// -----
+
+// CHECK: succeededParentOneOf
+func @succeededParentOneOf() {
+  "test.parent"() ({
+    "test.child_with_parent_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+   }) : () -> ()
+  return
+}
+
+// -----
+
+// CHECK: succeededParent1OneOf
+func @succeededParent1OneOf() {
+  "test.parent1"() ({
+    "test.child_with_parent_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+   }) : () -> ()
+  return
+}
+
+// -----
+
+func @failedParentOneOf_wrong_parent1() {
+  "some.otherop"() ({
+    // expected-error at +1 {{'test.child_with_parent_one_of' op expects parent op to be one of 'test.parent, test.parent1'}}
+    "test.child_with_parent_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+   }) : () -> ()
+}
+
+
 // -----
 
 func @failedSingleBlockImplicitTerminator_empty_block() {

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 8e5b380dff45..a4d74ebfd82b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -439,10 +439,18 @@ def BroadcastableOp : TEST_Op<"broadcastable", [ResultsBroadcastableShape]> {
   let results = (outs AnyTensor);
 }
 
-// There the "HasParent" trait.
-def ParentOp : TEST_Op<"parent">;
+// HasParent trait
+def ParentOp : TEST_Op<"parent"> {
+    let regions = (region AnyRegion);
+}
 def ChildOp : TEST_Op<"child", [HasParent<"ParentOp">]>;
 
+// ParentOneOf trait
+def ParentOp1 : TEST_Op<"parent1"> {
+  let regions = (region AnyRegion);
+}
+def ChildWithParentOneOf : TEST_Op<"child_with_parent_one_of",
+                                [ParentOneOf<["ParentOp", "ParentOp1"]>]>;
 
 def TerminatorOp : TEST_Op<"finish", [Terminator]>;
 def SingleBlockImplicitTerminatorOp : TEST_Op<"SingleBlockImplicitTerminator",