Add DialectType and generate docs for dialect types

Add new `typeDescription` (description was already used by base constraint class) field to type to allow writing longer descriptions about a type being defined. This allows for providing additional information/rationale for a defined type. This currently uses `description` as the heading/name for the type in the generated documentation.

PiperOrigin-RevId: 273299332

mlir/include/mlir/IR/OpBase.td

@@ -264,11 +264,21 @@ class Dialect {
 
 // A type, carries type constraints.
 class Type<Pred condition, string descr = ""> :
-    TypeConstraint<condition, descr>;
+    TypeConstraint<condition, descr> {
+  string typeDescription = "";
+}
 
 // Allows providing an alternative name and description to an existing type def.
 class TypeAlias<Type t, string description = t.description> :
-    Type<t.predicate, description>;
+    Type<t.predicate, description> {
+  let typeDescription = t.typeDescription;
+}
+
+// A type of a specific dialect.
+class DialectType<Dialect d, Pred condition, string descr = ""> :
+    Type<condition, descr> {
+  Dialect dialect = d;
+}
 
 // A variadic type constraint. It expands to zero or more of the base type. This
 // class is used for supporting variadic operands/results. An op can declare no

mlir/include/mlir/TableGen/Dialect.h

@@ -33,7 +33,7 @@ namespace tblgen {
 // and provides helper methods for accessing them.
 class Dialect {
 public:
-  explicit Dialect(const llvm::Record *def) : def(*def) {}
+  explicit Dialect(const llvm::Record *def) : def(def) {}
 
   // Returns the name of this dialect.
   StringRef getName() const;
@@ -55,8 +55,11 @@ public:
   // Compares two dialects by comparing the names of the dialects.
   bool operator<(const Dialect &other) const;
 
+  // Returns whether the dialect is defined.
+  operator bool() const { return def != nullptr; }
+
 private:
-  const llvm::Record &def;
+  const llvm::Record *def;
 };
 } // end namespace tblgen
 } // end namespace mlir

mlir/include/mlir/TableGen/Type.h

@@ -24,6 +24,7 @@
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/TableGen/Constraint.h"
+#include "mlir/TableGen/Dialect.h"
 
 namespace llvm {
 class DefInit;
@@ -46,6 +47,18 @@ public:
   bool isVariadic() const;
 };
 
+// Wrapper class with helper methods for accessing Types defined in TableGen.
+class Type : public TypeConstraint {
+public:
+  explicit Type(const llvm::Record *record);
+
+  // Returns the description of the type.
+  StringRef getTypeDescription() const;
+
+  // Returns the dialect for the type if defined.
+  Dialect getDialect() const;
+};
+
 } // end namespace tblgen
 } // end namespace mlir
 

mlir/lib/TableGen/Dialect.cpp

@@ -26,11 +26,11 @@ namespace mlir {
 namespace tblgen {
 
 StringRef tblgen::Dialect::getName() const {
-  return def.getValueAsString("name");
+  return def->getValueAsString("name");
 }
 
 StringRef tblgen::Dialect::getCppNamespace() const {
-  return def.getValueAsString("cppNamespace");
+  return def->getValueAsString("cppNamespace");
 }
 
 static StringRef getAsStringOrEmpty(const llvm::Record &record,
@@ -44,15 +44,15 @@ static StringRef getAsStringOrEmpty(const llvm::Record &record,
 }
 
 StringRef tblgen::Dialect::getSummary() const {
-  return getAsStringOrEmpty(def, "summary");
+  return getAsStringOrEmpty(*def, "summary");
 }
 
 StringRef tblgen::Dialect::getDescription() const {
-  return getAsStringOrEmpty(def, "description");
+  return getAsStringOrEmpty(*def, "description");
 }
 
 bool Dialect::operator==(const Dialect &other) const {
-  return &def == &other.def;
+  return def == other.def;
 }
 
 bool Dialect::operator<(const Dialect &other) const {

mlir/lib/TableGen/Type.cpp

@@ -23,16 +23,27 @@
 #include "llvm/TableGen/Record.h"
 
 using namespace mlir;
+using namespace mlir::tblgen;
 
-tblgen::TypeConstraint::TypeConstraint(const llvm::Record *record)
+TypeConstraint::TypeConstraint(const llvm::Record *record)
     : Constraint(Constraint::CK_Type, record) {
   assert(def->isSubClassOf("TypeConstraint") &&
          "must be subclass of TableGen 'TypeConstraint' class");
 }
 
-tblgen::TypeConstraint::TypeConstraint(const llvm::DefInit *init)
+TypeConstraint::TypeConstraint(const llvm::DefInit *init)
     : TypeConstraint(init->getDef()) {}
 
-bool tblgen::TypeConstraint::isVariadic() const {
+bool TypeConstraint::isVariadic() const {
   return def->isSubClassOf("Variadic");
 }
+
+Type::Type(const llvm::Record *record) : TypeConstraint(record) {}
+
+StringRef Type::getTypeDescription() const {
+  return def->getValueAsString("typeDescription");
+}
+
+Dialect Type::getDialect() const {
+  return Dialect(def->getValueAsDef("dialect"));
+}

mlir/tools/mlir-tblgen/OpDocGen.cpp

@@ -89,15 +89,23 @@ static void emitIfNotEmpty(StringRef str, raw_ostream &os) {
 
 static void emitOpDocForDialect(const Dialect &dialect,
                                 const std::vector<Operator> &ops,
+                                const std::vector<Type> &types,
                                 raw_ostream &os) {
   os << "# Dialect '" << dialect.getName() << "' definition\n";
   emitIfNotEmpty(dialect.getSummary(), os);
   emitIfNotEmpty(dialect.getDescription(), os);
 
-  // TODO(antiagainst): Add docs for types used (maybe dialect specific ones?)
-  // and link between use and def.
+  // TODO(antiagainst): Add link between use and def for types
+  if (!types.empty())
+    os << "## Type definition\n";
+  for (auto type : types) {
+    os << "### " << type.getDescription() << "\n";
+    emitDescription(type.getTypeDescription(), os);
+    os << "\n";
+  }
 
-  os << "## Operation definition\n";
+  if (!ops.empty())
+    os << "## Operation definition\n";
   for (auto op : ops) {
     os << "### " << op.getOperationName() << " (" << op.getQualCppClassName()
        << ")";
@@ -152,16 +160,25 @@ static void emitOpDocForDialect(const Dialect &dialect,
 }
 
 static void emitOpDoc(const RecordKeeper &recordKeeper, raw_ostream &os) {
-  const auto &defs = recordKeeper.getAllDerivedDefinitions("Op");
-  os << "<!-- Autogenerated by mlir-tblgen; don't manually edit -->\n";
+  const auto &opDefs = recordKeeper.getAllDerivedDefinitions("Op");
+  const auto &typeDefs = recordKeeper.getAllDerivedDefinitions("DialectType");
 
   std::map<Dialect, std::vector<Operator>> dialectOps;
-  for (auto *def : defs) {
-    Operator op(def);
+  std::map<Dialect, std::vector<Type>> dialectTypes;
+  for (auto *opDef : opDefs) {
+    Operator op(opDef);
     dialectOps[op.getDialect()].push_back(op);
   }
+  for (auto *typeDef : typeDefs) {
+    Type type(typeDef);
+    if (auto dialect = type.getDialect())
+      dialectTypes[dialect].push_back(type);
+  }
+
+  os << "<!-- Autogenerated by mlir-tblgen; don't manually edit -->\n";
   for (auto dialectWithOps : dialectOps)
-    emitOpDocForDialect(dialectWithOps.first, dialectWithOps.second, os);
+    emitOpDocForDialect(dialectWithOps.first, dialectWithOps.second,
+                        dialectTypes[dialectWithOps.first], os);
 }
 
 static mlir::GenRegistration