[mlir][Interfaces] Add a extraSharedClassDeclaration field

This field allows for defining a code block that is placed in both the interface
and trait declarations. This is very useful when defining a set of utilities to
expose on both the Interface class and the derived attribute/operation/type.

In non-static methods, `$_attr`/`$_op`/`$_type` (depending on the type of
interface) may be used to refer to an instance of the IR entity. In the interface
declaration, this is an instance of the interface class. In the trait declaration,
this is an instance of the concrete entity class (e.g. `IntegerAttr`, `FuncOp`, etc.).

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

Author: River707

mlir/docs/Interfaces.md

@@ -393,6 +393,18 @@ comprised of the following components:
         These declarations are _not_ implicitly visible in default
         implementations of interface methods, but static declarations may be
         accessed with full name qualification.
+*   Extra Shared Class Declarations (Optional: `extraSharedClassDeclaration`)
+    -   Additional C++ code that is injected into the declarations of both the
+        interface and trait class. This allows for defining methods and more
+        that are exposed on both the interface and trait class, e.g. to inject
+        utilties on both the interface and the derived entity implementing the
+        interface (e.g. attribute, operation, etc.).
+    -   In non-static methods, `$_attr`/`$_op`/`$_type`
+        (depending on the type of interface) may be used to refer to an
+        instance of the IR entity. In the interface declaration, the type of
+        the instance is the interface class. In the trait declaration, the
+        type of the instance is the concrete entity class
+        (e.g. `IntegerAttr`, `FuncOp`, etc.).
 
 `OpInterface` classes may additionally contain the following:
 

mlir/include/mlir/IR/OpBase.td

@@ -2226,6 +2226,10 @@ class Interface<string name> {
   // An optional code block containing extra declarations to place in the
   // interface declaration.
   code extraClassDeclaration = "";
+
+  // An optional code block containing extra declarations to place in both
+  // the interface and trait declaration.
+  code extraSharedClassDeclaration = "";
 }
 
 // AttrInterface represents an interface registered to an attribute.

mlir/include/mlir/TableGen/Interfaces.h

@@ -91,6 +91,10 @@ class Interface {
   // Return the traits extra class declaration code.
   llvm::Optional<StringRef> getExtraTraitClassDeclaration() const;
 
+  // Return the extra class declaration code shared between the interface and
+  // trait classes.
+  llvm::Optional<StringRef> getExtraSharedClassDeclaration() const;
+
   // Return the verify method body if it has one.
   llvm::Optional<StringRef> getVerify() const;
 

mlir/lib/TableGen/Interfaces.cpp

@@ -110,6 +110,12 @@ llvm::Optional<StringRef> Interface::getExtraTraitClassDeclaration() const {
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }
 
+// Return the shared extra class declaration code.
+llvm::Optional<StringRef> Interface::getExtraSharedClassDeclaration() const {
+  auto value = def->getValueAsString("extraSharedClassDeclaration");
+  return value.empty() ? llvm::Optional<StringRef>() : value;
+}
+
 // Return the body for this method if it has one.
 llvm::Optional<StringRef> Interface::getVerify() const {
   // Only OpInterface supports the verify method.

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

@@ -3,6 +3,24 @@
 
 include "mlir/IR/OpBase.td"
 
+def ExtraShardDeclsInterface : OpInterface<"ExtraShardDeclsInterface"> {
+  let extraSharedClassDeclaration = [{
+    bool sharedMethodDeclaration() {
+      return $_op.someOtherMethod();
+    }
+  }];
+}
+
+// DECL: class ExtraShardDeclsInterface
+// DECL:      bool sharedMethodDeclaration() {
+// DECL-NEXT:   return (*this).someOtherMethod();
+// DECL-NEXT: }
+
+// DECL: struct ExtraShardDeclsInterfaceTrait
+// DECL:      bool sharedMethodDeclaration() {
+// DECL-NEXT:   return (*static_cast<ConcreteOp *>(this)).someOtherMethod();
+// DECL-NEXT: }
+
 def TestOpInterface : OpInterface<"TestOpInterface"> {
   let description = [{some op interface description}];
 

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

@@ -104,6 +104,7 @@ class InterfaceGenerator {
   /// The format context to use for methods.
   tblgen::FmtContext nonStaticMethodFmt;
   tblgen::FmtContext traitMethodFmt;
+  tblgen::FmtContext extraDeclsFmt;
 };
 
 /// A specialized generator for attribute interfaces.
@@ -118,6 +119,7 @@ struct AttrInterfaceGenerator : public InterfaceGenerator {
     nonStaticMethodFmt.addSubst("_attr", castCode).withSelf(castCode);
     traitMethodFmt.addSubst("_attr",
                             "(*static_cast<const ConcreteAttr *>(this))");
+    extraDeclsFmt.addSubst("_attr", "(*this)");
   }
 };
 /// A specialized generator for operation interfaces.
@@ -132,6 +134,7 @@ struct OpInterfaceGenerator : public InterfaceGenerator {
         .withOp(castCode)
         .withSelf(castCode);
     traitMethodFmt.withOp("(*static_cast<ConcreteOp *>(this))");
+    extraDeclsFmt.withOp("(*this)");
   }
 };
 /// A specialized generator for type interfaces.
@@ -146,6 +149,7 @@ struct TypeInterfaceGenerator : public InterfaceGenerator {
     nonStaticMethodFmt.addSubst("_type", castCode).withSelf(castCode);
     traitMethodFmt.addSubst("_type",
                             "(*static_cast<const ConcreteType *>(this))");
+    extraDeclsFmt.addSubst("_type", "(*this)");
   }
 };
 } // namespace
@@ -415,6 +419,8 @@ void InterfaceGenerator::emitTraitDecl(Interface &interface,
   }
   if (auto extraTraitDecls = interface.getExtraTraitClassDeclaration())
     os << tblgen::tgfmt(*extraTraitDecls, &traitMethodFmt) << "\n";
+  if (auto extraTraitDecls = interface.getExtraSharedClassDeclaration())
+    os << tblgen::tgfmt(*extraTraitDecls, &traitMethodFmt) << "\n";
 
   os << "  };\n";
 }
@@ -469,6 +475,9 @@ void InterfaceGenerator::emitInterfaceDecl(Interface interface) {
   // Emit any extra declarations.
   if (Optional<StringRef> extraDecls = interface.getExtraClassDeclaration())
     os << *extraDecls << "\n";
+  if (Optional<StringRef> extraDecls =
+          interface.getExtraSharedClassDeclaration())
+    os << tblgen::tgfmt(*extraDecls, &extraDeclsFmt);
 
   os << "};\n";