[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

228

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

1064

?

mlir/include/mlir/TableGen/AttrOrTypeDef.h

118

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

118

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

[Mogball]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

160–161

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

233–236

Can you add the extra definitions here instead?

[bhatuzdaname]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

233–236

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

233–236

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

233–236

done, thank you!

[Mogball]

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

mlir/include/mlir/TableGen/Class.h

557

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

235

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;
 
   /// 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/include/mlir/TableGen/Class.h

public:
 
   /// Write the extra class definitions.
   void writeDefTo(raw_indented_ostream &os,
                   StringRef namePrefix) const override;
 
 private:
   /// The string of the extra class declarations. It is re-indented before
   /// printed.
   StringRef extraClassDeclaration;
   /// The string of the extra class definitions. It is re-indented before
   /// printed.
-  StringRef extraClassDefinition;
+  std::string extraClassDefinition;
 };
 
 /// A class used to emit C++ classes from Tablegen.  Contains a list of public
 /// methods and a list of private fields to be emitted.
 class Class {
 public:
   virtual ~Class() = default;
 

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 G_BuilderWithReturnTypeAttr : TestAttr<"BuilderWithReturnType"> {
   let parameters = (ins "int":$a);
   let genVerifyDecl = 1;
   let builders = [AttrBuilder<(ins), [{ return {}; }], "::mlir::Attribute">];
 }
 
 // DECL-LABEL: class BuilderWithReturnTypeAttr
 // DECL: ::mlir::Attribute get(
 // DECL: ::mlir::Attribute getChecked(
+
+def H_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:
   /// 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;
 };
 } // 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.
   for (const AttrOrTypeParameter &param : def.getParameters())
     if (param.isAnonymous())
for (auto &trait : def.getTraits()) {
     defParent.addTemplateParam(
         isa<NativeTrait>(&trait)
             ? cast<NativeTrait>(&trait)->getFullyQualifiedTraitName()
             : cast<InterfaceTrait>(&trait)->getFullyQualifiedTraitName());
   }
   defCls.addParent(std::move(defParent));
 }
 
+/// Extra class definitions have a `$cppClass` substitution that is to be
+/// replaced by the C++ class name.
+static std::string formatExtraDefinitions(const AttrOrTypeDef &def) {
+  if (Optional<StringRef> extraDef = def.getExtraDefs()) {
+    FmtContext ctx = FmtContext().addSubst("cppClass", def.getCppClassName());
+    return tgfmt(*extraDef, &ctx).str();
+  }
+  return "";
+}
+
 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);
+  Optional<StringRef> extraDecl = def.getExtraDecls();
+  std::string extraDef = formatExtraDefinitions(def);
+  defCls.declare<ExtraClassDeclaration>(extraDecl ? *extraDecl : "",
+                                        std::move(extraDef));
 }
 
 void DefGen::emitBuilders() {
   if (!def.skipDefaultBuilders()) {
     emitDefaultBuilder();
     if (def.genVerifyDecl())
       emitCheckedBuilder();
   }