[MLIR][ODS] Allow operations to specify interfaces using the HasParent trait constraint

mlir/include/mlir/IR/OpBase.td

@@ -141,6 +141,11 @@ class ParentOneOf<list<string> ops>
     : ParamNativeOpTrait<"HasParent", !interleave(ops, ", ")>,
       StructuralOpTrait;
 
+// Op's parent operation implements the provided interface.
+class HasInterfaceParent<string interface>
+    : ParamNativeOpTrait<"HasInterfaceParent", interface>,
+      StructuralOpTrait;
+
 // 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.

mlir/include/mlir/IR/OpDefinition.h

@@ -1334,6 +1334,24 @@ struct HasParent {
   };
 };
 
+/// This class provides a verifier for ops that are expecting their parent to
+/// implement a specific interface.
+template <typename ParentInterfaceType>
+struct HasInterfaceParent {
+  template <typename ConcreteType>
+  class Impl : public TraitBase<ConcreteType, Impl> {
+  public:
+    static LogicalResult verifyTrait(Operation *op) {
+      if (llvm::isa_and_nonnull<ParentInterfaceType>(op->getParentOp()))
+        return success();
+
+      return op->emitOpError()
+             << "expects parent op to implement interface '"
+             << ParentInterfaceType::getInterfaceName() << "'";
+    }
+  };
+};
+
 /// A trait for operations that have an attribute specifying operand segments.
 ///
 /// Certain operations can have multiple variadic operands and their size

mlir/test/IR/traits.mlir

@@ -668,3 +668,21 @@ func.func @failed_attr_traits() {
   "test.attr_with_trait"() {attr = 42 : i32} : () -> ()
   return
 }
+
+// -----
+
+func.func @has_interface_parent_trait() {
+  // CHECK: "test.interface_parent"() ({
+  // CHECK:    "test.interface_child"() : () -> ()
+  "test.interface_parent"() ({
+    "test.interface_child"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+func.func @illegal_interface_parent_trait() {
+  // expected-error@+1 {{'test.interface_child' op expects parent op to implement interface 'TestInterfaceParentInterface'}}
+  "test.interface_child"() : () -> ()
+}

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

@@ -147,4 +147,8 @@ def TestOptionallyImplementedOpInterface
   }];
 }
 
+def TestInterfaceParentInterface : OpInterface<"TestInterfaceParentInterface"> {
+    let cppNamespace = "::mlir";
+}
+
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

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

@@ -750,6 +750,14 @@ def ParentOp1 : TEST_Op<"parent1"> {
 def ChildWithParentOneOf : TEST_Op<"child_with_parent_one_of",
                                 [ParentOneOf<["ParentOp", "ParentOp1"]>]>;
 
+// HasInterfaceParent trait
+def InterfaceParentOp : TEST_Op<"interface_parent",
+                [SingleBlock, NoTerminator, TestInterfaceParentInterface]> {
+  let regions = (region SizedRegion<1>:$region);
+}
+def InterfaceChildOp : TEST_Op<"interface_child",
+                               [HasInterfaceParent<"mlir::TestInterfaceParentInterface">]>;
+
 def TerminatorOp : TEST_Op<"finish", [Terminator]>;
 def SingleBlockImplicitTerminatorOp : TEST_Op<"SingleBlockImplicitTerminator",
     [SingleBlockImplicitTerminator<"TerminatorOp">]> {

mlir/test/mlir-tblgen/op-interface.td

@@ -167,12 +167,17 @@ def DeclareMethodsOp : Op<TestDialect, "declare_methods_op",
 def DeclareMethodsWithDefaultOp : Op<TestDialect, "declare_methods_op",
       [DeclareOpInterfaceMethods<TestOpInterface, ["default_foo"]>]>;
 
+
 // DECL-LABEL: TestOpInterfaceInterfaceTraits
 // DECL: class TestOpInterface : public ::mlir::OpInterface<TestOpInterface, detail::TestOpInterfaceInterfaceTraits>
 
+// DECL: /// Returns the name of this interface.
+// DECL: static ::llvm::StringLiteral getInterfaceName() { return ::llvm::StringLiteral( "TestOpInterface"); }
+
 // DECL: /// some function comment
 // DECL: int foo(int input);
 
+
 // DECL-LABEL: struct TestOpInterfaceVerifyTrait
 // DECL: verifyTrait
 

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

@@ -494,6 +494,16 @@ void InterfaceGenerator::emitTraitDecl(const Interface &interface,
   os << "  };\n";
 }
 
+static void emitInterfaceNameGetter(const Interface &interface,
+                                    raw_ostream &os) {
+  if (!isa<OpInterface>(interface))
+    return;
+  os << "  /// Returns the name of this interface.\n"
+     << "  static ::llvm::StringLiteral getInterfaceName() { return "
+        "::llvm::StringLiteral( \""
+     << interface.getName() << "\"); }\n";
+}
+
 static void emitInterfaceDeclMethods(const Interface &interface,
                                      raw_ostream &os, StringRef valueType,
                                      bool isOpInterface,
@@ -553,6 +563,9 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
                       "  struct Trait : public detail::{0}Trait<{1}> {{};\n",
                       interfaceName, valueTemplate);
 
+  // Emit the name of the interface.
+  emitInterfaceNameGetter(interface, os);
+
   // Insert the method declarations.
   bool isOpInterface = isa<OpInterface>(interface);
   emitInterfaceDeclMethods(interface, os, valueType, isOpInterface,
@@ -588,7 +601,8 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
        << "    auto* interface = getInterfaceFor(base);\n"
        << "    if (!interface)\n"
           "      return false;\n"
-          "    " << interfaceName << " odsInterfaceInstance(base, interface);\n"
+          "    "
+       << interfaceName << " odsInterfaceInstance(base, interface);\n"
        << "    " << tblgen::tgfmt(extraClassOf->trim(), &extraClassOfFmt)
        << "\n  }\n";
   }