[mlir][docgen] Enable op grouping mechanism.

[jpienaar]

Summary

Add ability to be able to group ops together in documentation generation. This
results in section of ops grouped together under one summary & description.
This doesn't do anything special for the groups beyond nesting yet.

[scotttodd]

mlir/include/mlir/IR/OpBase.td

2529–2534

Did you consider adding a "group" or "docGroup" property to each op instead / in addition to this?

// One-line human-readable description of what the op does.
string summary = "";

// Additional, longer human-readable description of what the op does.
string description = "";

// Optional group label used by documentation generation.
string group = "";

That way, each op could say which group it belongs to, rather than need to keep a separate list of which ops belong to which group.

mlir/test/mlir-tblgen/gen-dialect-doc.td

21–25

Does this support multiple groups? If so, can you include multiple groups in this test?

[Mogball]

mlir/include/mlir/IR/OpBase.td

2531

Please add some comments on what the fields are (self explanatory but it always helps)

[jpienaar]

mlir/include/mlir/IR/OpBase.td

2529–2534

I didn't want to mix the op definitions with a documentation mechanism. E.g., the group they are in is purely "cosmetic" at this point. I could see a benefit though in that then one could use tablegen classes to create groups (if you already have a binary base class).

I'll think a bit.

[scotttodd]

mlir/include/mlir/IR/OpBase.td

2529–2534

Isn't the "description" similarly "cosmetic"?

Concrete example:
https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

//===----------------------------------------------------------------------===//
// TOSA Spec Section 2.4
// Operator Class: Elementwise unary/binary/ternary operators.
// Operator Subclass: Elementwise binary ops.
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// Operator: add
//===----------------------------------------------------------------------===//
def Tosa_AddOp : Tosa_ElemWiseBinaryOp<"add", [Commutative]> {
  let summary = "Elementwise addition operator";

  let description = [{
    Elementwise addition of input1 and input2. Axis of size 1 will be broadcast,
    as necessary.
  }];

Should a group be added like this:

def : OpDocGroup {
  let summary = "Elementwise binary ops";
  let description = "Elementwise binary ops (TOSA Spec Section 2.4).";
  let ops = [Tosa_AddOp, Tosa_ArithmeticRightShiftOp, Tosa_BitwiseAndOp, ...];
}

Or like this?

def Tosa_AddOp : Tosa_ElemWiseBinaryOp<"add", [Commutative]> {
  let summary = "Elementwise addition operator";

  let description = [{
    Elementwise addition of input1 and input2. Axis of size 1 will be broadcast,
    as necessary.
  }];

  let group = "elementwise_binary";

Maybe a combination of the two? (One for a label, another for the information backing that label)

[jpienaar]

mlir/include/mlir/IR/OpBase.td

2529–2534

Good question, I see the description as useful locally when one reads the file, while the group is more rendered side.

I think the two step is easier for folks to use though [I'll update post meetings :)]

[jpienaar]

Updating to refer to grouping with op.

[scotttodd]

Patched and tested downstream in IREE with a few groups in one of our larger dialects. Looks good! I'll need to tune our table of contents (TOC) depth level for each page, but that's to be expected.

mlir/test/mlir-tblgen/gen-dialect-doc.td

33–36

nit: in the spirit of https://llvm.org/docs/CodingStandards.html#namespace-indentation , I'd suggest putting a comment where the group is closed. It's less important in this small test, but larger files with 100s of lines between the group open and group close would benefit from that style being used.

[jpienaar]

In D153306#4438178, @scotttodd wrote:

Patched and tested downstream in IREE with a few groups in one of our larger dialects. Looks good! I'll need to tune our table of contents (TOC) depth level for each page, but that's to be expected.

I was actually thinking about that and how to handle ops not in a group is tricky. I could for example make all operand sub headings all be at the same depth, that would give (at max depth 2)

test.A
Goup
  test.B
test.C

rather than (depth 1)

test.A
Group
test.C

or (if one did as here and didnt' nest even ungrouped ops)

test.A
  Operands
Group
  test.b
test.C
  Operands

I could make it such that if you use any op grouping all ops get grouped and there is like a remainder section (I didn't like this one). Biasedly I looked at where this was used similarly and found that current approach covers it.

[scotttodd]

mlir/tools/mlir-tblgen/OpDocGen.cpp

371

I was actually thinking about that and how to handle ops not in a group is tricky.

I'd expect groups with one item to be nested as usual.

I have a practical example (with some screenshots) here: https://github.com/openxla/iree/pull/14194#discussion_r1239127215

I might explicitly put ungrouped ops into an "other" group, since it's ambiguous whether the common infrastructure should do that on its own, as you point out.

mlir/include/mlir/IR/OpBase.td

 }
 class Arg<Constraint constraint, string desc = "",
           list<OpVariableDecorator> decorators = []>
   : OpVariable<constraint, desc, decorators>;
 class Res<Constraint constraint, string desc = "",
           list<OpVariableDecorator> decorators = []>
   : OpVariable<constraint, desc, decorators>;
 
+// Marker to group ops together for documentation purposes.
+class OpDocGroup {
+  // Single line summary of the group of ops.
+  string summary;
+
+  // Longer description of documentation group.
+  string description;
+}
+
 // Base class for all ops.
 class Op<Dialect dialect, string mnemonic, list<Trait> props = []> {
   // The dialect of the op.
   Dialect opDialect = dialect;
 
   // The mnemonic of the op.
   string opName = mnemonic;
 
   // The C++ namespace to use for this op.
   string cppNamespace = dialect.cppNamespace;
 
   // One-line human-readable description of what the op does.
   string summary = "";
 
   // Additional, longer human-readable description of what the op does.
   string description = "";
 
+  // Optional. The group of ops this op is part of.
+  OpDocGroup opDocGroup = ?;
+
   // Dag containing the arguments of the op. Default to 0 arguments.
   dag arguments = (ins);
 
   // The list of results of the op. Default to 0 results.
   dag results = (outs);
 
   // The list of regions of the op. Default to 0 regions.
   dag regions = (region);
class Op<Dialect dialect, string mnemonic, list<Trait> props = []> {
   code extraClassDeclaration = ?;
 
   // Additional code that will be added to the generated source file. The
   // generated code is placed inside the op's C++ namespace. `$cppClass` is
   // replaced by the op's C++ class name.
   code extraClassDefinition = ?;
 }
 
 // The arguments of an op.
 class Arguments<dag args> {
   dag arguments = args;
 }
 
 // The results of an op.
 class Results<dag rets> {
   dag results = rets;
 }
 
 //===----------------------------------------------------------------------===//
 // Common op type constraints
 //===----------------------------------------------------------------------===//
 

mlir/test/mlir-tblgen/gen-dialect-doc.td

 // RUN: mlir-tblgen -gen-dialect-doc -I %S/../../include -dialect=test %s | FileCheck %s
 // RUN: mlir-tblgen -gen-dialect-doc -I %S/../../include -dialect=test_toc %s | FileCheck %s --check-prefix=CHECK_TOC
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
   let summary = "Dialect of ops to test";
   let description = [{
     Dialect without a [TOC] here.
     TOC added by tool.
   }];
   let cppNamespace = "NS";
 }
-def AOp : Op<Test_Dialect, "a", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
+
+def OpGroupA : OpDocGroup {
+  let summary = "Group of ops";
+  let description = "Grouped for some reason.";
+}
+
+let opDocGroup = OpGroupA in {
+def ADOp : Op<Test_Dialect, "d", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
+def AAOp : Op<Test_Dialect, "a", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
+}
+
+def OpGroupB : OpDocGroup {
+  let summary = "Other group of ops";
+  let description = "Grouped for some other reason.";
+}
+
+let opDocGroup = OpGroupB in {
+def ACOp : Op<Test_Dialect, "c", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
+def ABOp : Op<Test_Dialect, "b", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
+}
+
+def AEOp : Op<Test_Dialect, "e", [NoMemoryEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
 
 def TestAttr : DialectAttr<Test_Dialect, CPred<"true">> {
   let summary = "attribute summary";
   let description = "attribute description";
 }
 
 def TestType : DialectType<Test_Dialect, CPred<"true">> {
   let summary = "type summary";
def TestTypeDefParams : TypeDef<Test_Dialect, "TestTypeDefParams"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
 // CHECK: Dialect without a [TOC] here.
 // CHECK: TOC added by tool.
 // CHECK: [TOC]
 
 // CHECK-NOT: [TOC]
+// CHECK: test.e
+// CHECK: Group of ops
+// CHECK: test.a
+// CHECK: test.d
+// CHECK: Other group
+// CHECK: test.b
+// CHECK: test.c
 // CHECK: Traits: SingleBlockImplicitTerminator<YieldOp>
 // CHECK: Interfaces: NoMemoryEffect (MemoryEffectOpInterface)
 // CHECK: Effects: MemoryEffects::Effect{}
 
 // CHECK: ## Attribute constraint definition
 // CHECK: ### attribute summary
 // CHECK: attribute description
 

mlir/tools/mlir-tblgen/OpDocGen.cpp

 #include "llvm/TableGen/TableGenBackend.h"
 
 #include <set>
 #include <string>
 
 //===----------------------------------------------------------------------===//
 // Commandline Options
 //===----------------------------------------------------------------------===//
-static llvm::cl::OptionCategory docCat("Options for -gen-(attrdef|typedef|op|dialect)-doc");
-llvm::cl::opt<std::string> stripPrefix(
-    "strip-prefix",
+static llvm::cl::OptionCategory
+    docCat("Options for -gen-(attrdef|typedef|op|dialect)-doc");
+llvm::cl::opt<std::string>
+    stripPrefix("strip-prefix",
                 llvm::cl::desc("Strip prefix of the fully qualified names"),
                 llvm::cl::init("::mlir::"), llvm::cl::cat(docCat));
 
 using namespace llvm;
 using namespace mlir;
 using namespace mlir::tblgen;
 using mlir::tblgen::Operator;
 
 void mlir::tblgen::emitSummary(StringRef summary, raw_ostream &os) {
   if (!summary.empty()) {
static void emitAttrOrTypeDefAssemblyFormat(const AttrOrTypeDef &def,
   char prefix = isa<AttrDef>(def) ? '#' : '!';
   if (parameters.empty()) {
     os << "\nSyntax: `" << prefix << def.getDialect().getName() << "."
        << def.getMnemonic() << "`\n";
     return;
   }
 
   os << "\nSyntax:\n\n```\n"
-     << prefix << def.getDialect().getName() << "." << def.getMnemonic() << "<\n";
+     << prefix << def.getDialect().getName() << "." << def.getMnemonic()
+     << "<\n";
   for (const auto &it : llvm::enumerate(parameters)) {
     const AttrOrTypeParameter &param = it.value();
     os << "  " << param.getSyntax();
     if (it.index() < (parameters.size() - 1))
       os << ",";
     os << "   # " << param.getName() << "\n";
   }
   os << ">\n```\n";
static void emitAttrOrTypeDefDoc(const RecordKeeper &recordKeeper,
   for (const llvm::Record *def : defs)
     emitAttrOrTypeDefDoc(AttrOrTypeDef(def), os);
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Documentation
 //===----------------------------------------------------------------------===//
 
-static void emitDialectDoc(const Dialect &dialect,
-                           ArrayRef<Attribute> attributes,
-                           ArrayRef<AttrDef> attrDefs, ArrayRef<Operator> ops,
-                           ArrayRef<Type> types, ArrayRef<TypeDef> typeDefs,
-                           raw_ostream &os) {
-  os << "# '" << dialect.getName() << "' Dialect\n\n";
-  emitIfNotEmpty(dialect.getSummary(), os);
-  emitIfNotEmpty(dialect.getDescription(), os);
+struct OpDocGroup {
+  const Dialect &getDialect() const { return ops.front().getDialect(); }
 
-  // Generate a TOC marker except if description already contains one.
-  llvm::Regex r("^[[:space:]]*\\[TOC\\]$", llvm::Regex::RegexFlags::Newline);
-  if (!r.match(dialect.getDescription()))
-    os << "[TOC]\n\n";
+  // Returns the summary description of the section.
+  std::string summary = "";
+
+  // Returns the description of the section.
+  StringRef description = "";
+
+  // Instances inside the section.
+  std::vector<Operator> ops;
+};
 
+static void maybeNest(bool nest, llvm::function_ref<void(raw_ostream &os)> fn,
+                      raw_ostream &os) {
+  std::string str;
+  llvm::raw_string_ostream ss(str);
+  fn(ss);
+  for (StringRef x : llvm::split(ss.str(), "\n")) {
+    if (nest && x.starts_with("#"))
+      os << "#";
+    os << x << "\n";
+  }
+}
+
+static void emitBlock(ArrayRef<Attribute> attributes,
+                      ArrayRef<AttrDef> attrDefs, ArrayRef<OpDocGroup> ops,
+                      ArrayRef<Type> types, ArrayRef<TypeDef> typeDefs,
+                      raw_ostream &os) {
   if (!ops.empty()) {
     os << "## Operation definition\n\n";
-    for (const Operator &op : ops)
+    for (const OpDocGroup &grouping : ops) {
+      bool nested = grouping.ops.size() > 1;
+      maybeNest(
+          nested,
+          [&](raw_ostream &os) {
+            if (nested) {
+              os << "## " << StringRef(grouping.summary).trim() << "\n\n";
+              emitDescription(grouping.description, os);
+              os << "\n\n";
+            }
+            for (const Operator &op : grouping.ops) {
               emitOpDoc(op, os);
             }
+          },
+          os);
+    }
+  }
 
   if (!attributes.empty()) {
     os << "## Attribute constraint definition\n\n";
     for (const Attribute &attr : attributes)
       emitAttrDoc(attr, os);
   }
 
   if (!attrDefs.empty()) {
static void emitBlock(ArrayRef<Attribute> attributes,
 
   if (!typeDefs.empty()) {
     os << "## Type definition\n\n";
     for (const TypeDef &def : typeDefs)
       emitAttrOrTypeDefDoc(def, os);
   }
 }
 
+static void emitDialectDoc(const Dialect &dialect,
+                           ArrayRef<Attribute> attributes,
+                           ArrayRef<AttrDef> attrDefs, ArrayRef<OpDocGroup> ops,
+                           ArrayRef<Type> types, ArrayRef<TypeDef> typeDefs,
+                           raw_ostream &os) {
+  os << "# '" << dialect.getName() << "' Dialect\n\n";
+  emitIfNotEmpty(dialect.getSummary(), os);
+  emitIfNotEmpty(dialect.getDescription(), os);
+
+  // Generate a TOC marker except if description already contains one.
+  llvm::Regex r("^[[:space:]]*\\[TOC\\]$", llvm::Regex::RegexFlags::Newline);
+  if (!r.match(dialect.getDescription()))
+    os << "[TOC]\n\n";
+
+  emitBlock(attributes, attrDefs, ops, types, typeDefs, os);
+}
+
 static bool emitDialectDoc(const RecordKeeper &recordKeeper, raw_ostream &os) {
   std::vector<Record *> dialectDefs =
       recordKeeper.getAllDerivedDefinitionsIfDefined("Dialect");
   SmallVector<Dialect> dialects(dialectDefs.begin(), dialectDefs.end());
   std::optional<Dialect> dialect = findDialectToGenerate(dialects);
   if (!dialect)
     return true;
 
   std::vector<Record *> opDefs = getRequestedOpDefinitions(recordKeeper);
   std::vector<Record *> attrDefs =
       recordKeeper.getAllDerivedDefinitionsIfDefined("DialectAttr");
   std::vector<Record *> typeDefs =
       recordKeeper.getAllDerivedDefinitionsIfDefined("DialectType");
   std::vector<Record *> typeDefDefs =
       recordKeeper.getAllDerivedDefinitionsIfDefined("TypeDef");
   std::vector<Record *> attrDefDefs =
       recordKeeper.getAllDerivedDefinitionsIfDefined("AttrDef");
 
   std::vector<Attribute> dialectAttrs;
   std::vector<AttrDef> dialectAttrDefs;
-  std::vector<Operator> dialectOps;
+  std::vector<OpDocGroup> dialectOps;
   std::vector<Type> dialectTypes;
   std::vector<TypeDef> dialectTypeDefs;
+
   llvm::SmallDenseSet<Record *> seen;
   auto addIfInDialect = [&](llvm::Record *record, const auto &def, auto &vec) {
-    if (seen.insert(record).second && def.getDialect() == *dialect)
+    if (seen.insert(record).second && def.getDialect() == *dialect) {
       vec.push_back(def);
+      return true;
+    }
+    return false;
   };
 
+  SmallDenseMap<Record *, OpDocGroup> opDocGroup;
+
   for (Record *def : attrDefDefs)
     addIfInDialect(def, AttrDef(def), dialectAttrDefs);
   for (Record *def : attrDefs)
     addIfInDialect(def, Attribute(def), dialectAttrs);
-  for (Record *def : opDefs)
-    addIfInDialect(def, Operator(def), dialectOps);
+  for (Record *def : opDefs) {
+    if (Record *group = def->getValueAsOptionalDef("opDocGroup")) {
+      OpDocGroup &op = opDocGroup[group];
+      addIfInDialect(def, Operator(def), op.ops);
+    } else {
+      OpDocGroup op;
+      op.ops.emplace_back(def);
+      addIfInDialect(def, op, dialectOps);
+    }
+  }
+  for (Record *rec :
+       recordKeeper.getAllDerivedDefinitionsIfDefined("OpDocGroup")) {
+    if (opDocGroup[rec].ops.empty())
+      continue;
+    opDocGroup[rec].summary = rec->getValueAsString("summary");
+    opDocGroup[rec].description = rec->getValueAsString("description");
+    dialectOps.push_back(opDocGroup[rec]);
+  }
   for (Record *def : typeDefDefs)
     addIfInDialect(def, TypeDef(def), dialectTypeDefs);
   for (Record *def : typeDefs)
     addIfInDialect(def, Type(def), dialectTypes);
 
+  // Sort alphabetically ignorning dialect for ops and section name for
+  // sections.
+  // TODO: The sorting order could be revised, currently attempting to sort of
+  // keep in alphabetical order.
+  std::sort(dialectOps.begin(), dialectOps.end(),
+            [](const OpDocGroup &lhs, const OpDocGroup &rhs) {
+              auto getDesc = [](const OpDocGroup &arg) -> StringRef {
+                if (!arg.summary.empty())
+                  return arg.summary;
+                return arg.ops.front().getDef().getValueAsString("opName");
+              };
+              return getDesc(lhs).compare_insensitive(getDesc(rhs)) < 0;
+            });
+
   os << "<!-- Autogenerated by mlir-tblgen; don't manually edit -->\n";
   emitDialectDoc(*dialect, dialectAttrs, dialectAttrDefs, dialectOps,
                  dialectTypes, dialectTypeDefs, os);
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Gen Registration