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

[jurahul]

Summary

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

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1162

Does << llvm::makeArrayRef(...) not work?

[jurahul]

Address review comments

[jurahul]

mlir/include/mlir/IR/OpDefinition.h

1162

Looks like it does. I have adopted that now.

[jurahul]

The 2 test failures look unrelated to this change. As such, I'd appreciate if someone can help land this change.

mlir/include/mlir/IR/OpBase.td

 // Op's regions have a single block with the specified terminator.
 class SingleBlockImplicitTerminator<string op>
     : ParamNativeOpTrait<"SingleBlockImplicitTerminator", op>;
 
 // Op's parent operation is the provided one.
 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.
 def FirstAttrDerivedResultType :
   GenInternalOpTrait<"FirstAttrDerivedResultType">;
 
 // TODO(antiagainst): Turn the following into normal traits and generate
 // verification for them.

mlir/include/mlir/IR/OpDefinition.h

public:
     Block *getBody(unsigned idx = 0) {
       Region &region = this->getOperation()->getRegion(idx);
       assert(!region.empty() && "unexpected empty region");
       return &region.front();
     }
   };
 };
 
-/// 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()...})
+             << "'";
     }
   };
 };
 
 /// A trait for operations that have an attribute specifying operand segments.
 ///
 /// Certain operations can have multiple variadic operands and their size
 /// relationship is not always known statically. For such cases, we need

mlir/test/IR/traits.mlir

func @failedHasParent_wrong_parent() {
   "some.op"() ({
    // expected-error@+1 {{'test.child' op expects parent op 'test.parent'}}
     "test.child"() : () -> ()
   }) : () -> ()
 }
 
 // -----
 
+// 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@+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() {
    // expected-error@+1 {{'test.SingleBlockImplicitTerminator' op expects a non-empty block}}
   "test.SingleBlockImplicitTerminator"() ({
   ^entry:
   }) : () -> ()
 }
 
 // -----

mlir/test/lib/Dialect/Test/TestOps.td

def IfFirstOperandIsNoneThenSoIsSecond :
   let arguments = (ins AnyType:$x, AnyType:$y);
 }
 
 def BroadcastableOp : TEST_Op<"broadcastable", [ResultsBroadcastableShape]> {
   let arguments = (ins Variadic<AnyTensor>);
   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",
     [SingleBlockImplicitTerminator<"TerminatorOp">]> {
   let regions = (region SizedRegion<1>:$region);
 }
 
 def I32ElementsAttrOp : TEST_Op<"i32ElementsAttr"> {