[mlir] Restrict dialect doc gen to a single dialect

[rriddle]

Summary

In the overwhelmingly majority of cases only one dialect is generated at a time
anyways, and this restriction more easily catches user error when multiple
dialects might be generated. We hit this semi-recently with the PDL dialect,
and circt+other downstream users are also actively hitting this as well.

[lattner]

Nice, thank you for fixing this!

[jpienaar]

Thanks, and yes with Marius 's change for doc includes (couple months back :-)) there isn't really reason to have multiple be generated website side. Being explicit here is better.

Did none of the cmake files need to change? If I'm reading this change correctly it just bails out if no dialect is specified, and so I'd expect some to fail (well except if the default cmake config already does set this flag)

[rriddle]

In D125651#3514675, @jpienaar wrote:

Thanks, and yes with Marius 's change for doc includes (couple months back :-)) there isn't really reason to have multiple be generated website side. Being explicit here is better.

Did none of the cmake files need to change? If I'm reading this change correctly it just bails out if no dialect is specified, and so I'd expect some to fail (well except if the default cmake config already does set this flag)

It bails if no dialect is specified and there isn't a single dialect (which matches the dialectgen handling). We realistically could always enforce the flag, but I'd need to update uses of dialect gen as well.

[rriddle]

In D125651#3514675, @jpienaar wrote:

Did none of the cmake files need to change? If I'm reading this change correctly it just bails out if no dialect is specified, and so I'd expect some to fail (well except if the default cmake config already does set this flag)

I had to update a couple, which conveniently looked alright on the website only because the main dialect was alphabetically ordered before the one that got included. I'll try to see if in a follow up we can roll this into add_mlir_dialect, to handle all of this automatically.

mlir/include/mlir/Dialect/AMX/CMakeLists.txt

 add_mlir_dialect(AMX amx)
-add_mlir_doc(AMX AMX Dialects/ -gen-dialect-doc)
+add_mlir_doc(AMX AMX Dialects/ -gen-dialect-doc -dialect=amx)
 
 set(LLVM_TARGET_DEFINITIONS AMX.td)
 mlir_tablegen(AMXConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRAMXConversionsIncGen)

mlir/include/mlir/Dialect/ArmNeon/CMakeLists.txt

 add_mlir_dialect(ArmNeon arm_neon)
-add_mlir_doc(ArmNeon ArmNeon Dialects/ -gen-dialect-doc)
+add_mlir_doc(ArmNeon ArmNeon Dialects/ -gen-dialect-doc -dialect=arm_neon)
 
 set(LLVM_TARGET_DEFINITIONS ArmNeon.td)
 mlir_tablegen(ArmNeonConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRArmNeonConversionsIncGen)

mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt

 add_mlir_dialect(ArmSVE arm_sve ArmSVE)
-add_mlir_doc(ArmSVE ArmSVE Dialects/ -gen-dialect-doc)
+add_mlir_doc(ArmSVE ArmSVE Dialects/ -gen-dialect-doc -dialect=arm_sve)
 
 set(LLVM_TARGET_DEFINITIONS ArmSVE.td)
 mlir_tablegen(ArmSVEConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRArmSVEConversionsIncGen)

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

-// RUN: mlir-tblgen -gen-dialect-doc -I %S/../../include %s | FileCheck %s
+// 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/Interfaces/SideEffectInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
   let summary = "Dialect of ops to test";
   let description = [{
let description = [{
     [TOC]
 
     here.
   }];
   let cppNamespace = "NS";
 }
 def BOp : Op<Toc_Dialect, "b", []>;
 
-// CHECK: Dialect with
-// CHECK: [TOC]
-// CHECK: here.
+// CHECK_TOC: Dialect with
+// CHECK_TOC: [TOC]
+// CHECK_TOC: here.

mlir/tools/mlir-tblgen/DialectGen.cpp

 //===- DialectGen.cpp - MLIR dialect definitions generator ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // DialectGen uses the description of dialects to generate C++ definitions.
 //
 //===----------------------------------------------------------------------===//
 
+#include "DialectGenUtilities.h"
 #include "mlir/TableGen/Class.h"
 #include "mlir/TableGen/CodeGenHelpers.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Interfaces.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/Trait.h"
 #include "llvm/ADT/Optional.h"

 filterForDialect(ArrayRef<llvm::Record *> records, Dialect &dialect) {
   auto filterFn = [&](const llvm::Record *record) {
     return T(record).getDialect() == dialect;
   };
   return {DialectFilterIterator(records.begin(), records.end(), filterFn),
           DialectFilterIterator(records.end(), records.end(), filterFn)};
 }
 
-static Optional<Dialect>
-findSelectedDialect(ArrayRef<const llvm::Record *> dialectDefs) {
-  // Select the dialect to gen for.
-  if (dialectDefs.size() == 1 && selectedDialect.getNumOccurrences() == 0) {
-    return Dialect(dialectDefs.front());
+Optional<Dialect> tblgen::findDialectToGenerate(ArrayRef<Dialect> dialects) {
+  if (dialects.empty()) {
+    llvm::errs() << "no dialect was found\n";
+    return llvm::None;
   }
 
+  // Select the dialect to gen for.
+  if (dialects.size() == 1 && selectedDialect.getNumOccurrences() == 0)
+    return dialects.front();
+
   if (selectedDialect.getNumOccurrences() == 0) {
     llvm::errs() << "when more than 1 dialect is present, one must be selected "
                     "via '-dialect'\n";
     return llvm::None;
   }
 
-  const auto *dialectIt =
-      llvm::find_if(dialectDefs, [](const llvm::Record *def) {
-        return Dialect(def).getName() == selectedDialect;
-      });
-  if (dialectIt == dialectDefs.end()) {
+  const auto *dialectIt = llvm::find_if(dialects, [](const Dialect &dialect) {
+    return dialect.getName() == selectedDialect;
+  if (dialectIt == dialects.end()) {
     llvm::errs() << "selected dialect with '-dialect' does not exist\n";
     return llvm::None;
   }
-  return Dialect(*dialectIt);
+  return *dialectIt;
 }
 
 //===----------------------------------------------------------------------===//
 // GEN: Dialect declarations
 //===----------------------------------------------------------------------===//
 
 /// The code block for the start of a dialect class declaration.
 ///

 static bool emitDialectDecls(const llvm::RecordKeeper &recordKeeper,
                              raw_ostream &os) {
   emitSourceFileHeader("Dialect Declarations", os);
 
   auto dialectDefs = recordKeeper.getAllDerivedDefinitions("Dialect");
   if (dialectDefs.empty())
     return false;
 
-  Optional<Dialect> dialect = findSelectedDialect(dialectDefs);
+  SmallVector<Dialect> dialects(dialectDefs.begin(), dialectDefs.end());
+  Optional<Dialect> dialect = findDialectToGenerate(dialects);
   if (!dialect)
     return true;
   auto attrDefs = recordKeeper.getAllDerivedDefinitions("DialectAttr");
   auto typeDefs = recordKeeper.getAllDerivedDefinitions("DialectType");
   emitDialectDecl(*dialect, filterForDialect<Attribute>(attrDefs, *dialect),
                   filterForDialect<Type>(typeDefs, *dialect), os);
   return false;
 }

 static bool emitDialectDefs(const llvm::RecordKeeper &recordKeeper,
                             raw_ostream &os) {
   emitSourceFileHeader("Dialect Definitions", os);
 
   auto dialectDefs = recordKeeper.getAllDerivedDefinitions("Dialect");
   if (dialectDefs.empty())
     return false;
 
-  Optional<Dialect> dialect = findSelectedDialect(dialectDefs);
+  SmallVector<Dialect> dialects(dialectDefs.begin(), dialectDefs.end());
+  Optional<Dialect> dialect = findDialectToGenerate(dialects);
   if (!dialect)
     return true;
   emitDialectDef(*dialect, os);
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // GEN: Dialect registration hooks

mlir/tools/mlir-tblgen/DialectGenUtilities.h

+//===- DialectGenUtilities.h - Utilities for dialect generation -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIRTBLGEN_DIALECTGENUTILITIES_H_
+#define MLIR_TOOLS_MLIRTBLGEN_DIALECTGENUTILITIES_H_
+
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+namespace tblgen {
+class Dialect;
+
+/// Find the dialect selected by the user to generate for. Returns None if no
+/// dialect was found, or if more than one potential dialect was found.
+Optional<Dialect> findDialectToGenerate(ArrayRef<Dialect> dialects);
+} // namespace tblgen
+} // namespace mlir
+
+#endif // MLIR_TOOLS_MLIRTBLGEN_DIALECTGENUTILITIES_H_

mlir/tools/mlir-tblgen/OpDocGen.cpp

 //===- OpDocGen.cpp - MLIR operation documentation generator --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // OpDocGen uses the description of operations to generate documentation for the
 // operations.
 //
 //===----------------------------------------------------------------------===//
 
+#include "DialectGenUtilities.h"
 #include "DocGenUtilities.h"
 #include "OpGenHelpers.h"
 #include "mlir/Support/IndentedOstream.h"
 #include "mlir/TableGen/AttrOrTypeDef.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Operator.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 #include "llvm/TableGen/TableGenBackend.h"
 
 #include <set>
 
 using namespace llvm;
 using namespace mlir;
 using namespace mlir::tblgen;
 
 using mlir::tblgen::Operator;
 
-extern llvm::cl::opt<std::string> selectedDialect;
-
 // Emit the description by aligning the text to the left per line (e.g.,
 // removing the minimum indentation across the block).
 //
 // This expects that the description in the tablegen file is already formatted
 // in a way the user wanted but has some additional indenting due to being
 // nested in the op definition.
 void mlir::tblgen::emitDescription(StringRef description, raw_ostream &os) {
   raw_indented_ostream ros(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) {
-  if (selectedDialect.getNumOccurrences() &&
-      dialect.getName() != selectedDialect)
-    return;
   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";
static void emitDialectDoc(const Dialect &dialect,
 
   if (!typeDefs.empty()) {
     os << "## Type definition\n\n";
     for (const TypeDef &def : typeDefs)
       emitAttrOrTypeDefDoc(def, os);
   }
 }
 
-static void emitDialectDoc(const RecordKeeper &recordKeeper, raw_ostream &os) {
+static bool emitDialectDoc(const RecordKeeper &recordKeeper, raw_ostream &os) {
   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::set<Dialect> dialectsWithDocs;
+  llvm::SetVector<Dialect, SmallVector<Dialect>, std::set<Dialect>>
+      dialectsWithDocs;
 
   llvm::StringMap<std::vector<Attribute>> dialectAttrs;
   llvm::StringMap<std::vector<AttrDef>> dialectAttrDefs;
   llvm::StringMap<std::vector<Operator>> dialectOps;
   llvm::StringMap<std::vector<Type>> dialectTypes;
   llvm::StringMap<std::vector<TypeDef>> dialectTypeDefs;
   for (Record *attrDef : attrDefs) {
     Attribute attr(attrDef);
for (Record *typeDef : typeDefs) {
     }
   }
   for (Record *typeDef : typeDefDefs) {
     TypeDef type(typeDef);
     dialectTypeDefs[type.getDialect().getName()].push_back(type);
     dialectsWithDocs.insert(type.getDialect());
   }
 
+  Optional<Dialect> dialect =
+      findDialectToGenerate(dialectsWithDocs.getArrayRef());
+  if (!dialect)
+    return true;
+
   os << "<!-- Autogenerated by mlir-tblgen; don't manually edit -->\n";
-  for (const Dialect &dialect : dialectsWithDocs) {
-    StringRef dialectName = dialect.getName();
-    emitDialectDoc(dialect, dialectAttrs[dialectName],
-                   dialectAttrDefs[dialectName], dialectOps[dialectName],
-                   dialectTypes[dialectName], dialectTypeDefs[dialectName], os);
-  }
+  StringRef dialectName = dialect->getName();
+  emitDialectDoc(*dialect, dialectAttrs[dialectName],
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Gen Registration
 //===----------------------------------------------------------------------===//
 
 static mlir::GenRegistration
     genAttrRegister("gen-attrdef-doc",
genTypeRegister("gen-typedef-doc", "Generate dialect type documentation",
                     [](const RecordKeeper &records, raw_ostream &os) {
                       emitAttrOrTypeDefDoc(records, os, "TypeDef");
                       return false;
                     });
 
 static mlir::GenRegistration
     genRegister("gen-dialect-doc", "Generate dialect documentation",
                 [](const RecordKeeper &records, raw_ostream &os) {
-                  emitDialectDoc(records, os);
-                  return false;
+                  return emitDialectDoc(records, os);
                 });