From ec0b53d4e49a6efdcdb1b911944236e997f00a91 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Wed, 10 Nov 2021 16:09:43 -0800
Subject: [PATCH] [mlir] Add traits, interfaces, effects to generated docs

Simply emit traits, interfaces & effects (with some minimal formatting) to the
generated docs to make this information easier to find in the docs.

Differential Revision: https://reviews.llvm.org/D113539
---
 mlir/test/mlir-tblgen/gen-dialect-doc.td |  6 ++-
 mlir/tools/mlir-tblgen/OpDocGen.cpp      | 61 +++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/mlir/test/mlir-tblgen/gen-dialect-doc.td b/mlir/test/mlir-tblgen/gen-dialect-doc.td
index a49ebaae8800..1edb423cee05 100644
--- a/mlir/test/mlir-tblgen/gen-dialect-doc.td
+++ b/mlir/test/mlir-tblgen/gen-dialect-doc.td
@@ -1,6 +1,7 @@
 // RUN: mlir-tblgen -gen-dialect-doc -I %S/../../include %s | FileCheck %s
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
@@ -11,12 +12,15 @@ def Test_Dialect : Dialect {
   }];
   let cppNamespace = "NS";
 }
-def AOp : Op<Test_Dialect, "a", []>;
+def AOp : Op<Test_Dialect, "a", [NoSideEffect, SingleBlockImplicitTerminator<"YieldOp">]>;
 
 // CHECK: Dialect without a [TOC] here.
 // CHECK: TOC added by tool.
 // CHECK: [TOC]
 // CHECK-NOT: [TOC]
+// CHECK: Traits: SingleBlockImplicitTerminator<YieldOp>
+// CHECK: Interfaces: NoSideEffect (MemoryEffectOpInterface)
+// CHECK: Effects: MemoryEffects::Effect{}
 
 def Toc_Dialect : Dialect {
   let name = "test_toc";
diff --git a/mlir/tools/mlir-tblgen/OpDocGen.cpp b/mlir/tools/mlir-tblgen/OpDocGen.cpp
index ab3248390ea8..db73bbb6e6e2 100644
--- a/mlir/tools/mlir-tblgen/OpDocGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDocGen.cpp
@@ -60,7 +60,7 @@ static void emitIfNotEmpty(StringRef str, raw_ostream &os) {
 template <typename T>
 static void emitNamedConstraint(const T &it, raw_ostream &os) {
   if (!it.name.empty())
-    os << "`" << it.name << "`";
+    os << "| `" << it.name << "`";
   else
     os << "&laquo;unnamed&raquo;";
   os << " | " << it.constraint.getSummary() << "\n";
@@ -88,6 +88,61 @@ static void emitAssemblyFormat(StringRef opName, StringRef format,
   os << "```\n\n";
 }
 
+static void emitOpTraitsDoc(Operator op, raw_ostream &os) {
+  // TODO: We should link to the trait/documentation of it. That also means we
+  // should add descriptions to traits that can be queried.
+  // Collect using set to sort effects, interfaces & traits.
+  std::set<std::string> effects, interfaces, traits;
+  for (auto &trait : op.getTraits()) {
+    if (auto *nt = dyn_cast<PredTrait>(&trait))
+      continue;
+
+    std::string name = trait.getDef().getName().str();
+    StringRef ref = name;
+    StringRef traitName = trait.getDef().getValueAsString("trait");
+    traitName.consume_back("::Trait");
+    traitName.consume_back("::Impl");
+    if (ref.startswith("anonymous_"))
+      name = traitName.str();
+    if (auto *nt = dyn_cast<InterfaceTrait>(&trait)) {
+      if (trait.getDef().isSubClassOf("SideEffectsTraitBase")) {
+        auto effectName = trait.getDef().getValueAsString("baseEffectName");
+        effectName.consume_front("::");
+        effectName.consume_front("mlir::");
+        std::string effectStr;
+        llvm::raw_string_ostream os(effectStr);
+        os << effectName << "{";
+        auto list = trait.getDef().getValueAsListOfDefs("effects");
+        llvm::interleaveComma(list, os, [&](Record *rec) {
+          StringRef effect = rec->getValueAsString("effect");
+          effect.consume_front("::");
+          effect.consume_front("mlir::");
+          os << effect << " on " << rec->getValueAsString("resource");
+        });
+        os << "}";
+        effects.insert(os.str());
+        name.append(llvm::formatv(" ({0})", traitName).str());
+      }
+      interfaces.insert(name);
+      continue;
+    }
+
+    traits.insert(name);
+  }
+  if (!traits.empty()) {
+    llvm::interleaveComma(traits, os << "\nTraits: ");
+    os << "\n";
+  }
+  if (!interfaces.empty()) {
+    llvm::interleaveComma(interfaces, os << "\nInterfaces: ");
+    os << "\n";
+  }
+  if (!effects.empty()) {
+    llvm::interleaveComma(effects, os << "\nEffects: ");
+    os << "\n";
+  }
+}
+
 static void emitOpDoc(Operator op, raw_ostream &os) {
   os << llvm::formatv("### `{0}` ({1})\n", op.getOperationName(),
                       op.getQualCppClassName());
@@ -101,6 +156,8 @@ static void emitOpDoc(Operator op, raw_ostream &os) {
   if (op.hasDescription())
     mlir::tblgen::emitDescription(op.getDescription(), os);
 
+  emitOpTraitsDoc(op, os);
+
   // Emit attributes.
   if (op.getNumAttributes() != 0) {
     // TODO: Attributes are only documented by TableGen name, with no further
@@ -110,7 +167,7 @@ static void emitOpDoc(Operator op, raw_ostream &os) {
        << "| :-------: | :-------: | ----------- |\n";
     for (const auto &it : op.getAttributes()) {
       StringRef storageType = it.attr.getStorageType();
-      os << "`" << it.name << "` | " << storageType << " | "
+      os << "| `" << it.name << "` | " << storageType << " | "
          << it.attr.getSummary() << "\n";
     }
   }