[mlir] separable registration of operation interfaces

This is similar to attribute and type interfaces and mostly the same mechanism
(FallbackModel / ExternalModel, ODS generation). There are minor differences in
how the concept-based polymorphism is implemented for operations that are
accounted for by ODS backends, and this essentially adds a test and exposes the
API.

Reviewed By: rriddle

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

Author: ftynse

mlir/docs/Interfaces.md

@@ -207,12 +207,12 @@ if (ExampleOpInterface example = dyn_cast<ExampleOpInterface>(op))
   llvm::errs() << "hook returned = " << example.exampleInterfaceHook() << "\n";
 ```
 
-#### External Models for Attribute/Type Interfaces
+#### External Models for Attribute, Operation and Type Interfaces
 
-It may be desirable to provide an interface implementation for an attribute or a
-type without modifying the definition of said attribute or type. Notably, this
-allows to implement interfaces for attributes and types outside of the dialect
-that defines them and, in particular, provide interfaces for built-in types.
+It may be desirable to provide an interface implementation for an IR object
+without modifying the definition of said object. Notably, this allows to
+implement interfaces for attributes, operations and types outside of the dialect
+that defines them, for example, to provide interfaces for built-in types.
 
 This is achieved by extending the concept-based polymorphism model with two more
 classes derived from `Concept` as follows.
@@ -261,9 +261,9 @@ struct ExampleTypeInterfaceTraits {
 };
 ```
 
-External models can be provided for attribute and type interfaces by deriving
-either `FallbackModel` or `ExternalModel` and by registering the model class
-with the attribute or type class in a given context. Other contexts will not see
+External models can be provided for attribute, operation and type interfaces by
+deriving either `FallbackModel` or `ExternalModel` and by registering the model
+class with the relevant class in a given context. Other contexts will not see
 the interface unless registered.
 
 ```c++

mlir/include/mlir/IR/OpDefinition.h

@@ -1621,6 +1621,19 @@ class Op : public OpState, public Traits<ConcreteType>... {
         reinterpret_cast<Operation *>(const_cast<void *>(pointer)));
   }
 
+  /// Attach the given models as implementations of the corresponding interfaces
+  /// for the concrete operation.
+  template <typename... Models>
+  static void attachInterface(MLIRContext &context) {
+    AbstractOperation *abstract = AbstractOperation::lookupMutable(
+        ConcreteType::getOperationName(), &context);
+    if (!abstract)
+      llvm::report_fatal_error(
+          "Attempting to attach an interface to an unregistered operation " +
+          ConcreteType::getOperationName() + ".");
+    abstract->interfaceMap.insert<Models...>();
+  }
+
 private:
   /// Trait to check if T provides a 'fold' method for a single result op.
   template <typename T, typename... Args>

mlir/include/mlir/IR/OperationSupport.h

@@ -159,7 +159,9 @@ class AbstractOperation {
   /// Look up the specified operation in the specified MLIRContext and return a
   /// pointer to it if present.  Otherwise, return a null pointer.
   static const AbstractOperation *lookup(StringRef opName,
-                                         MLIRContext *context);
+                                         MLIRContext *context) {
+    return lookupMutable(opName, context);
+  }
 
   /// This constructor is used by Dialect objects when they register the list of
   /// operations they contain.
@@ -190,6 +192,15 @@ class AbstractOperation {
                     GetCanonicalizationPatternsFn getCanonicalizationPatterns,
                     detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait);
 
+  /// Give Op access to lookupMutable.
+  template <typename ConcreteType, template <typename T> class... Traits>
+  friend class Op;
+
+  /// Look up the specified operation in the specified MLIRContext and return a
+  /// pointer to it if present.  Otherwise, return a null pointer.
+  static AbstractOperation *lookupMutable(StringRef opName,
+                                          MLIRContext *context);
+
   /// A map of interfaces that were registered to this operation.
   detail::InterfaceMap interfaceMap;
 

mlir/lib/IR/MLIRContext.cpp

@@ -693,8 +693,8 @@ ParseResult AbstractOperation::parseAssembly(OpAsmParser &parser,
 
 /// Look up the specified operation in the operation set and return a pointer
 /// to it if present. Otherwise, return a null pointer.
-const AbstractOperation *AbstractOperation::lookup(StringRef opName,
-                                                   MLIRContext *context) {
+AbstractOperation *AbstractOperation::lookupMutable(StringRef opName,
+                                                    MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredOperations.find(opName);
   if (it != impl.registeredOperations.end())

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

@@ -88,6 +88,24 @@ def TestExternalAttrInterface : AttrInterface<"TestExternalAttrInterface"> {
   ];
 }
 
+def TestExternalOpInterface : OpInterface<"TestExternalOpInterface"> {
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<"Returns the length of the operation name plus arg.",
+      "unsigned", "getNameLengthPlusArg", (ins "unsigned":$arg)>,
+    StaticInterfaceMethod<
+      "Returns the length of the operation name plus arg twice.", "unsigned",
+      "getNameLengthPlusArgTwice", (ins "unsigned":$arg)>,
+    InterfaceMethod<
+      "Returns the length of the product of the operation name and arg.",
+      "unsigned", "getNameLengthTimesArg", (ins "unsigned":$arg), "",
+      "return arg * $_op->getName().getStringRef().size();">,
+    StaticInterfaceMethod<"Returns the length of the operation name minus arg.",
+      "unsigned", "getNameLengthMinusArg", (ins "unsigned":$arg), "",
+      "return ConcreteOp::getOperationName().size() - arg;">,
+  ];
+}
+
 def TestEffectOpInterface
     : EffectOpInterfaceBase<"TestEffectOpInterface",
                             "::mlir::TestEffects::Effect"> {

mlir/unittests/IR/InterfaceAttachmentTest.cpp

@@ -12,10 +12,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "gtest/gtest.h"
 
 #include "../../test/lib/Dialect/Test/TestAttributes.h"
+#include "../../test/lib/Dialect/Test/TestDialect.h"
 #include "../../test/lib/Dialect/Test/TestTypes.h"
 
 using namespace mlir;
@@ -150,4 +152,72 @@ TEST(InterfaceAttachment, Attribute) {
   EXPECT_EQ(iface.getSomeNumber(), 42);
 }
 
+/// External interface model for the module operation. Only provides non-default
+/// methods.
+struct TestExternalOpModel
+    : public TestExternalOpInterface::ExternalModel<TestExternalOpModel,
+                                                    ModuleOp> {
+  unsigned getNameLengthPlusArg(Operation *op, unsigned arg) const {
+    return op->getName().getStringRef().size() + arg;
+  }
+
+  static unsigned getNameLengthPlusArgTwice(unsigned arg) {
+    return ModuleOp::getOperationName().size() + 2 * arg;
+  }
+};
+
+/// External interface model for the func operation. Provides non-deafult and
+/// overrides default methods.
+struct TestExternalOpOverridingModel
+    : public TestExternalOpInterface::FallbackModel<
+          TestExternalOpOverridingModel> {
+  unsigned getNameLengthPlusArg(Operation *op, unsigned arg) const {
+    return op->getName().getStringRef().size() + arg;
+  }
+
+  static unsigned getNameLengthPlusArgTwice(unsigned arg) {
+    return FuncOp::getOperationName().size() + 2 * arg;
+  }
+
+  unsigned getNameLengthTimesArg(Operation *op, unsigned arg) const {
+    return 42;
+  }
+
+  static unsigned getNameLengthMinusArg(unsigned arg) { return 21; }
+};
+
+TEST(InterfaceAttachment, Operation) {
+  MLIRContext context;
+
+  // Initially, the operation doesn't have the interface.
+  auto moduleOp = ModuleOp::create(UnknownLoc::get(&context));
+  ASSERT_FALSE(isa<TestExternalOpInterface>(moduleOp.getOperation()));
+
+  // We can attach an external interface and now the operaiton has it.
+  ModuleOp::attachInterface<TestExternalOpModel>(context);
+  auto iface = dyn_cast<TestExternalOpInterface>(moduleOp.getOperation());
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getNameLengthPlusArg(10), 16u);
+  EXPECT_EQ(iface.getNameLengthTimesArg(3), 18u);
+  EXPECT_EQ(iface.getNameLengthPlusArgTwice(18), 42u);
+  EXPECT_EQ(iface.getNameLengthMinusArg(5), 1u);
+
+  // Default implementation can be overridden.
+  auto funcOp = FuncOp::create(UnknownLoc::get(&context), "function",
+                               FunctionType::get(&context, {}, {}));
+  ASSERT_FALSE(isa<TestExternalOpInterface>(funcOp.getOperation()));
+  FuncOp::attachInterface<TestExternalOpOverridingModel>(context);
+  iface = dyn_cast<TestExternalOpInterface>(funcOp.getOperation());
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getNameLengthPlusArg(10), 14u);
+  EXPECT_EQ(iface.getNameLengthTimesArg(0), 42u);
+  EXPECT_EQ(iface.getNameLengthPlusArgTwice(8), 20u);
+  EXPECT_EQ(iface.getNameLengthMinusArg(1000), 21u);
+
+  // Another context doesn't have the interfaces registered.
+  MLIRContext other;
+  auto otherModuleOp = ModuleOp::create(UnknownLoc::get(&other));
+  ASSERT_FALSE(isa<TestExternalOpInterface>(otherModuleOp.getOperation()));
+}
+
 } // end namespace