[mlir][tblgen] Add support for extraClassDefinition in AttrDef

[bhatuzdaname]

Summary

For AttrDef declarations, place specified code in extraClassDefinition into the generated *.cpp.inc file.

[rriddle]

Please add a test in attrdefs.td and add documentation to docs/AttributesAndTypes.md.

mlir/include/mlir/IR/AttrTypeBase.td

220

Use class instead of attribute, this tablegen class is for both attributes and types.

[bhatuzdaname]

addressed comment, added test, clang-formatted files.

[rriddle]

LG, thanks!

mlir/docs/AttributesAndTypes.md

1041

?

mlir/include/mlir/TableGen/AttrOrTypeDef.h

114

This looks unrelated, and I think this is also causing the debian builder to fail (it's failing on clang format)

[bhatuzdaname]

Fix clang-format issue, improve documentation

[bhatuzdaname]

Thanks for the review! Could you commit it for me? I do not have commit access.

mlir/include/mlir/TableGen/AttrOrTypeDef.h

114

Weird. Might've come from the clang-format on my system. Reverted it.

[Mogball]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

162–163

This doesn't need to be declared as a class field.

238

Can you add the extra definitions here instead?

[bhatuzdaname]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

238

I had initially done that but ran into some lifetime issue with the underlying string for the definition which resulted in printing of a lot of garbage in the generated .inc file. So that's why I moved it into the class definition. I'll look into this again, maybe I missed something then

[Mogball]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

238

The field is a StringRef so you can change that in ExtraClassDeclaration.

[bhatuzdaname]

Remove extraclass string field, add decl+def in a single go.

[bhatuzdaname]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

238

done, thank you!

[Mogball]

Two minor comments. Otherwise LG. Thanks for making this!

mlir/include/mlir/TableGen/Class.h

557 ↗

(On Diff #445398)

extraClassDeclaration doesn't need to be std::string since it has no format requirements and can reference the TableGen string safely

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

240

The string isn't needed afterwards and can be moved into the object

[bhatuzdaname]

Addressed additional comments. Thanks for the review! Could you please commit it for me?

mlir/docs/AttributesAndTypes.md

public:
   /// initialized yet.
   Type getBody() { return getImpl()->containedType; }
 
   /// Returns the name.
   StringRef getName() { return getImpl()->name; }
 };
 ```
 
 ### Extra declarations
 
 The declarative Attribute and Type definitions try to auto-generate as much
 logic and methods as possible. With that said, there will always be long-tail
-cases that won't be covered. For such cases, `extraClassDeclaration` can be
-used. Code within the `extraClassDeclaration` field will be copied literally to
-the generated C++ Attribute or Type class.
-
-Note that `extraClassDeclaration` is a mechanism intended for long-tail cases by
-power users; for not-yet-implemented widely-applicable cases, improving the
-infrastructure is preferable.
+cases that won't be covered. For such cases, `extraClassDeclaration` and
+`extraClassDefinition` can be used. Code within the `extraClassDeclaration`
+field will be copied literally to the generated C++ Attribute or Type class.
+Code within `extraClassDefinition` will be added to the generated source file
+inside the class's C++ namespace. The substitution `$cppClass` will be replaced
+by the Attribute or Type's C++ class name.
+
+Note that these are mechanisms intended for long-tail cases by power users; for
+not-yet-implemented widely-applicable cases, improving the infrastructure is
+preferable.
 
 ### Registering with the Dialect
 
 Once the attributes and types have been defined, they must then be registered
 with the parent `Dialect`. This is done via the `addAttributes` and `addTypes`
 methods. Note that when registering, the full definition of the storage classes
 must be visible.
 

mlir/include/mlir/IR/AttrTypeBase.td

class AttrOrTypeDef<string valueType, string name, list<Trait> defTraits,
   // be provided.
   bit skipDefaultBuilders = 0;
 
   // Generate the verify and getChecked methods.
   bit genVerifyDecl = 0;
 
   // Extra code to include in the class declaration.
   code extraClassDeclaration = [{}];
+
+  // Additional code that will be added to the generated source file. The
+  // generated code is placed inside the class's C++ namespace. `$cppClass` is
+  // replaced by the class name.
+  code extraClassDefinition = [{}];
 }
 
 // Define a new attribute, named `name`, belonging to `dialect` that inherits
 // from the given C++ base class.
 class AttrDef<Dialect dialect, string name, list<Trait> traits = [],
               string baseCppClass = "::mlir::Attribute">
     : DialectAttr<dialect, CPred<"">, /*descr*/"">,
       AttrOrTypeDef<"Attr", name, traits, baseCppClass> {

mlir/include/mlir/TableGen/AttrOrTypeDef.h

public:
   }
   bool operator!=(const AttrOrTypeParameter &other) const {
     return !(*this == other);
   }
 
 private:
   /// A parameter can be either a string or a def. Get a potentially null value
   /// from the def.
-  template <typename InitT>
-  auto getDefValue(StringRef name) const;
+  template <typename InitT> auto getDefValue(StringRef name) const;
 
   /// The underlying tablegen parameter list this parameter is a part of.
   const llvm::DagInit *def;
   /// The index of the parameter within the parameter list (`def`).
   unsigned index;
 };
 
 //===----------------------------------------------------------------------===//
public:
 
   /// Return true if we need to generate the verify declaration and getChecked
   /// method.
   bool genVerifyDecl() const;
 
   /// Returns the def's extra class declaration code.
   Optional<StringRef> getExtraDecls() const;
 
+  /// Returns the def's extra class definition code.
+  Optional<StringRef> getExtraDefs() const;
+
   /// Get the code location (for error printing).
   ArrayRef<SMLoc> getLoc() const;
 
   /// Returns true if the default get/getChecked methods should be skipped
   /// during generation.
   bool skipDefaultBuilders() const;
 
   /// Returns the builders of this def.

mlir/lib/TableGen/AttrOrTypeDef.cpp

bool AttrOrTypeDef::genVerifyDecl() const {
   return def->getValueAsBit("genVerifyDecl");
 }
 
 Optional<StringRef> AttrOrTypeDef::getExtraDecls() const {
   auto value = def->getValueAsString("extraClassDeclaration");
   return value.empty() ? Optional<StringRef>() : value;
 }
 
+Optional<StringRef> AttrOrTypeDef::getExtraDefs() const {
+  auto value = def->getValueAsString("extraClassDefinition");
+  return value.empty() ? Optional<StringRef>() : value;
+}
+
 ArrayRef<SMLoc> AttrOrTypeDef::getLoc() const { return def->getLoc(); }
 
 bool AttrOrTypeDef::skipDefaultBuilders() const {
   return def->getValueAsBit("skipDefaultBuilders");
 }
 
 bool AttrOrTypeDef::operator==(const AttrOrTypeDef &other) const {
   return def == other.def;

mlir/test/mlir-tblgen/attrdefs.td

def F_ParamWithAccessorTypeAttr : TestAttr<"ParamWithAccessorType"> {
   let parameters = (ins AttrParameter<"std::string", "", "StringRef">:$param);
 }
 
 // DECL-LABEL: class ParamWithAccessorTypeAttr
 // DECL: StringRef getParam()
 // DEF: ParamWithAccessorTypeAttrStorage
 // DEF: ParamWithAccessorTypeAttrStorage(std::string param)
 // DEF: StringRef ParamWithAccessorTypeAttr::getParam()
+
+def G_TestExtraClassAttr : TestAttr<"TestExtraClass"> {
+  let extraClassDeclaration = [{
+    /// Test Method
+    static int getFoo(int i);
+  }];
+  let extraClassDefinition = [{
+    int $cppClass::getFoo(int i) {
+      return i+1;
+    }
+  }];
+}
+
+// DECL-LABEL: TestExtraClassAttr : public ::mlir::Attribute
+// DECL: /// Test Method
+// DECL-NEXT: static int getFoo(int i);
+
+// DEF-LABEL: int TestExtraClassAttr::getFoo(int i) {
+// DEF: return i+1;
+// DEF-NEXT: }

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

private:
   /// Emit a verifier for the def.
   void emitVerifier();
   /// Emit parsers and printers.
   void emitParserPrinter();
   /// Emit parameter accessors, if required.
   void emitAccessors();
   /// Emit interface methods.
   void emitInterfaceMethods();
+  /// Emit extra class definition, if specified.
+  void emitExtraClassDefinition();
 
   //===--------------------------------------------------------------------===//
   // Builder Emission
 
   /// Emit the default builder `Attribute::get`
   void emitDefaultBuilder();
   /// Emit the checked builder `Attribute::getChecked`
   void emitCheckedBuilder();
private:
   /// An optional attribute or type storage class. The storage class will
   /// exist if and only if the def has more than zero parameters.
   Optional<Class> storageCls;
 
   /// The C++ base value of the def, either "Attribute" or "Type".
   StringRef valueType;
   /// The prefix/suffix of the TableGen def name, either "Attr" or "Type".
   StringRef defType;
+  /// Hand-written extra class definitions.
+  std::string extraClassDefinition;
 };
 } // namespace
 
 DefGen::DefGen(const AttrOrTypeDef &def)
     : def(def), params(def.getParameters()), defCls(def.getCppClassName()),
       valueType(isa<AttrDef>(def) ? "Attribute" : "Type"),
       defType(isa<AttrDef>(def) ? "Attr" : "Type") {
   // Check that all parameters have names.
DefGen::DefGen(const AttrOrTypeDef &def)
   // Emit the mnemonic, if there is one, and any associated parser and printer.
   if (def.getMnemonic())
     emitParserPrinter();
   // Emit accessors
   if (def.genAccessors())
     emitAccessors();
   // Emit trait interface methods
   emitInterfaceMethods();
+  // Emit hand written class definitions
+  emitExtraClassDefinition();
   defCls.finalize();
   // Emit a storage class if one is needed
   if (storageCls && def.genStorageClass())
     emitStorageClass();
 }
 
 void DefGen::createParentWithTraits() {
   ParentClass defParent(strfmt("::mlir::{0}::{1}Base", valueType, defType));
   defParent.addTemplateParam(def.getCppClassName());
   defParent.addTemplateParam(def.getCppBaseClassName());
   defParent.addTemplateParam(storageCls
                                  ? strfmt("{0}::{1}", def.getStorageNamespace(),
                                           def.getStorageClassName())
                                  : strfmt("::mlir::{0}Storage", valueType));
   for (auto &trait : def.getTraits()) {
     defParent.addTemplateParam(
         isa<NativeTrait>(&trait)
             ? cast<NativeTrait>(&trait)->getFullyQualifiedTraitName()
             : cast<InterfaceTrait>(&trait)->getFullyQualifiedTraitName());
   }
   defCls.addParent(std::move(defParent));
 }
 
+void DefGen::emitExtraClassDefinition() {
+  if (Optional<StringRef> extraDef = def.getExtraDefs()) {
+    FmtContext ctx = FmtContext().addSubst("cppClass", def.getCppClassName());
+    extraClassDefinition = tgfmt(*extraDef, &ctx).str();
+    defCls.declare<ExtraClassDeclaration>("", extraClassDefinition);
+  }
+}
+
 void DefGen::emitTopLevelDeclarations() {
   // Inherit constructors from the attribute or type class.
   defCls.declare<VisibilityDeclaration>(Visibility::Public);
   defCls.declare<UsingDeclaration>("Base::Base");
 
   // Emit the extra declarations first in case there's a definition in there.
   if (Optional<StringRef> extraDecl = def.getExtraDecls())
     defCls.declare<ExtraClassDeclaration>(*extraDecl);
 }
 
 void DefGen::emitBuilders() {
   if (!def.skipDefaultBuilders()) {
     emitDefaultBuilder();
     if (def.genVerifyDecl())
       emitCheckedBuilder();
   }
   for (auto &builder : def.getBuilders()) {
     emitCustomBuilder(builder);