[mlir][Interfaces] Add a extraSharedClassDeclaration field

[rriddle]

Summary

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.).

Fixes #53000

[mehdi_amini]

Do you have in-tree example of uses of this?

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

478

Did you get a warning about this when running arc?

[mehdi_amini]

In D116961#3232664, @mehdi_amini wrote:

Do you have in-tree example of uses of this?

Actually: I'll look at the new ShapedType interface.

[mehdi_amini]

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

478

Ping here?

[rriddle]

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

478

No I didn't, but I think my local setup is a bit wonky. Thanks for pointing this out.

[jpienaar]

mlir/docs/Interfaces.md

397

injected?

404

I read this and wonder if it is "this" or this :)

mlir/include/mlir/TableGen/Interfaces.h

94

Could you add what it is shared between?

mlir/docs/Interfaces.md

*   Methods (`methods`)
     -   The structure of these methods is defined below.
 *   Extra Class Declarations (Optional: `extraClassDeclaration`)
     -   Additional C++ code that is generated in the declaration of the
         interface class. This allows for defining methods and more on the user
         facing interface class, that do not need to hook into the IR entity.
         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:
 
 *   Verifier (`verify`)
     -   A C++ code block containing additional verification applied to the
         operation that the interface is attached to.
     -   The structure of this code block corresponds 1-1 with the structure of a
         [`Trait::verifyTrait`](Traits.md) method.

mlir/include/mlir/IR/OpBase.td

class Interface<string name> {
   string cppNamespace = "";
 
   // The list of methods defined by this interface.
   list<InterfaceMethod> methods = [];
 
   // 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.
 class AttrInterface<string name> : Interface<name>, InterfaceTrait<name>;
 
 // OpInterface represents an interface registered to an operation.
 class OpInterface<string name> : Interface<name>, OpInterfaceTrait<name>;
 

mlir/include/mlir/TableGen/Interfaces.h

public:
   llvm::Optional<StringRef> getDescription() const;
 
   // Return the interfaces extra class declaration code.
   llvm::Optional<StringRef> getExtraClassDeclaration() const;
 
   // 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;
 
   // Returns the Tablegen definition this interface was constructed from.
   const llvm::Record &getDef() const { return *def; }
 
 private:
   // The TableGen definition of this interface.

mlir/lib/TableGen/Interfaces.cpp

 }
 
 // Return the traits extra class declaration code.
 llvm::Optional<StringRef> Interface::getExtraTraitClassDeclaration() const {
   auto value = def->getValueAsString("extraTraitClassDeclaration");
   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.
   if (!isa<OpInterface>(this))
     return llvm::None;
   auto value = def->getValueAsString("verify");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }

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

 // RUN: mlir-tblgen -gen-op-interface-decls -I %S/../../include %s | FileCheck %s --check-prefix=DECL
 // RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck %s --check-prefix=OP_DECL
 
 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}];
 
   let methods = [
     InterfaceMethod<
       /*desc=*/[{some function comment}],
       /*retTy=*/"int",
       /*methodName=*/"foo",

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

protected:
   StringRef valueType;
   /// The C++ base interface type.
   StringRef interfaceBaseType;
   /// The name of the typename for the value template.
   StringRef valueTemplate;
   /// The format context to use for methods.
   tblgen::FmtContext nonStaticMethodFmt;
   tblgen::FmtContext traitMethodFmt;
+  tblgen::FmtContext extraDeclsFmt;
 };
 
 /// A specialized generator for attribute interfaces.
 struct AttrInterfaceGenerator : public InterfaceGenerator {
   AttrInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(records.getAllDerivedDefinitions("AttrInterface"),
                            os) {
     valueType = "::mlir::Attribute";
     interfaceBaseType = "AttributeInterface";
     valueTemplate = "ConcreteAttr";
     StringRef castCode = "(tablegen_opaque_val.cast<ConcreteAttr>())";
     nonStaticMethodFmt.addSubst("_attr", castCode).withSelf(castCode);
     traitMethodFmt.addSubst("_attr",
                             "(*static_cast<const ConcreteAttr *>(this))");
+    extraDeclsFmt.addSubst("_attr", "(*this)");
   }
 };
 /// A specialized generator for operation interfaces.
 struct OpInterfaceGenerator : public InterfaceGenerator {
   OpInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(getAllOpInterfaceDefinitions(records), os) {
     valueType = "::mlir::Operation *";
     interfaceBaseType = "OpInterface";
     valueTemplate = "ConcreteOp";
     StringRef castCode = "(llvm::cast<ConcreteOp>(tablegen_opaque_val))";
     nonStaticMethodFmt.addSubst("_this", "impl")
         .withOp(castCode)
         .withSelf(castCode);
     traitMethodFmt.withOp("(*static_cast<ConcreteOp *>(this))");
+    extraDeclsFmt.withOp("(*this)");
   }
 };
 /// A specialized generator for type interfaces.
 struct TypeInterfaceGenerator : public InterfaceGenerator {
   TypeInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(records.getAllDerivedDefinitions("TypeInterface"),
                            os) {
     valueType = "::mlir::Type";
     interfaceBaseType = "TypeInterface";
     valueTemplate = "ConcreteType";
     StringRef castCode = "(tablegen_opaque_val.cast<ConcreteType>())";
     nonStaticMethodFmt.addSubst("_type", castCode).withSelf(castCode);
     traitMethodFmt.addSubst("_type",
                             "(*static_cast<const ConcreteType *>(this))");
+    extraDeclsFmt.addSubst("_type", "(*this)");
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // GEN: Interface definitions
 //===----------------------------------------------------------------------===//
 
if (auto verify = interface.getVerify()) {
     tblgen::FmtContext verifyCtx;
     verifyCtx.withOp("op");
     os << "    static ::mlir::LogicalResult verifyTrait(::mlir::Operation *op) "
           "{\n      "
        << tblgen::tgfmt(verify->trim(), &verifyCtx) << "\n    }\n";
   }
   if (auto extraTraitDecls = interface.getExtraTraitClassDeclaration())
     os << tblgen::tgfmt(*extraTraitDecls, &traitMethodFmt) << "\n";
+  if (auto extraTraitDecls = interface.getExtraSharedClassDeclaration())
+    os << tblgen::tgfmt(*extraTraitDecls, &traitMethodFmt) << "\n";
 
   os << "  };\n";
 }
 
 void InterfaceGenerator::emitInterfaceDecl(Interface interface) {
   llvm::SmallVector<StringRef, 2> namespaces;
   llvm::SplitString(interface.getCppNamespace(), namespaces, "::");
   for (StringRef ns : namespaces)
for (auto &method : interface.getMethods()) {
     emitMethodNameAndArgs(method, os, valueType, /*addThisArg=*/false,
                           /*addConst=*/!isOpInterface);
     os << ";\n";
   }
 
   // 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";
 
   os << "namespace detail {\n";
   emitTraitDecl(interface, interfaceName, interfaceTraitsName);
   os << "}// namespace detail\n";
 
   emitModelMethodsDef(interface);