[mlir] separable registration of operation interfaces

[ftynse]

Summary

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.

[ftynse]

Extend the documentation.

[rriddle]

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

102

105

mlir/docs/Interfaces.md

 
 /// Later, we can query if a specific operation(like 'MyOp') overrides the given
 /// interface.
 Operation *op = ...;
 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.
 
 ```c++
 struct ExampleTypeInterfaceTraits {
   struct Concept {
     virtual unsigned exampleInterfaceHook(Type type) const = 0;
struct ExternalModel : public FallbackModel<ConcreteModel> {
     unsigned exampleInterfaceHook(Type type) const override {
       // Default implementation can be provided here.
       return type.cast<ConcreteType>().callSomeTypeSpecificMethod();
     }
   };
 };
 ```
 
-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++
 /// External interface implementation for a concrete class. This does not
 /// require modifying the definition of the type class itself.
 struct ExternalModelExample
     : public ExampleTypeInterface::ExternalModel<ExternalModelExample,
                                                  IntegerType> {

mlir/include/mlir/IR/OpDefinition.h

#endif
   const void *getAsOpaquePointer() const {
     return static_cast<const void *>((Operation *)*this);
   }
   static ConcreteOpType getFromOpaquePointer(const void *pointer) {
     return ConcreteOpType(
         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>
   using has_single_result_fold =
       decltype(std::declval<T>().fold(std::declval<ArrayRef<Attribute>>()));
   template <typename T>
   using detect_has_single_result_fold =
       llvm::is_detected<has_single_result_fold, T>;

mlir/include/mlir/IR/OperationSupport.h

public:
   }
 
   /// Returns true if the operation has a particular trait.
   bool hasTrait(TypeID traitID) const { return hasTraitFn(traitID); }
 
   /// 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.
   template <typename T>
   static void insert(Dialect &dialect) {
     insert(T::getOperationName(), dialect, TypeID::get<T>(),
            T::getParseAssemblyFn(), T::getPrintAssemblyFn(),
            T::getVerifyInvariantsFn(), T::getFoldHookFn(),
private:
   AbstractOperation(StringRef name, Dialect &dialect, TypeID typeID,
                     ParseAssemblyFn &&parseAssembly,
                     PrintAssemblyFn &&printAssembly,
                     VerifyInvariantsFn &&verifyInvariants,
                     FoldHookFn &&foldHook,
                     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;
 
   /// Internal callback hooks provided by the op implementation.
   FoldHookFn foldHookFn;
   GetCanonicalizationPatternsFn getCanonicalizationPatternsFn;
   HasTraitFn hasTraitFn;
   ParseAssemblyFn parseAssemblyFn;

mlir/lib/IR/MLIRContext.cpp

 
 ParseResult AbstractOperation::parseAssembly(OpAsmParser &parser,
                                              OperationState &result) const {
   return parseAssemblyFn(parser, result);
 }
 
 /// 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,
+AbstractOperation *AbstractOperation::lookupMutable(StringRef opName,
                                                     MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredOperations.find(opName);
   if (it != impl.registeredOperations.end())
     return &it->second;
   return nullptr;
 }
 
 void AbstractOperation::insert(

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

def TestExternalAttrInterface : AttrInterface<"TestExternalAttrInterface"> {
   let cppNamespace = "::mlir";
   let methods = [
     InterfaceMethod<"Gets the dialect pointer.", "const ::mlir::Dialect *",
       "getDialectPtr">,
     StaticInterfaceMethod<"Returns some number.", "int", "getSomeNumber">,
   ];
 }
 
+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"> {
   let cppNamespace = "::mlir";
 }
 
 class TestEffect<string effectName>
     : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
 
 class TestEffects<list<TestEffect> effects = []>
    : SideEffectsTraitBase<TestEffectOpInterface, effects>;
 
 def TestConcreteEffect : TestEffect<"TestEffects::Concrete">;
 
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

mlir/unittests/IR/InterfaceAttachmentTest.cpp

 //===- InterfaceAttachmentTest.cpp - Test attaching interfaces ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This implements the tests for attaching interfaces to attributes and types
 // without having to specify them on the attribute or type class directly.
 //
 //===----------------------------------------------------------------------===//
 
 #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;
 using namespace mlir::test;
 
 namespace {
 
 /// External interface model for the integer type. Only provides non-default
TEST(InterfaceAttachment, Attribute) {
   ASSERT_FALSE(attr.isa<TestExternalAttrInterface>());
   IntegerAttr::attachInterface<TextExternalIntegerAttrModel>(context);
   auto iface = attr.dyn_cast<TestExternalAttrInterface>();
   ASSERT_TRUE(iface != nullptr);
   EXPECT_EQ(iface.getDialectPtr(), &attr.getDialect());
   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