[mlir][tblgen] Refact mlir-tblgen main into its own library

This has previously been done for `mlir-opt` and `mlir-reduce` and roughly the same approach has been done here.

The use case for having a separate library is that it is easier for downstream to make custom TableGen backends/executable that work on top of the utilities that are defined in `mlir/TableGen`.
The customization point here is the same one as for any upstream TableGen backends: One can add a new generator by simply creating a global instance of `mlir::GenRegistration`.

Differential Revision: https://reviews.llvm.org/D131112

8 files changed, 225 insertions, 139 deletions

diff --git a/mlir/include/mlir/Tools/mlir-tblgen/MlirTblgenMain.h b/mlir/include/mlir/Tools/mlir-tblgen/MlirTblgenMain.h
new file mode 100644
index 000000000000..b824a9f1a0cf
--- /dev/null
+++ b/mlir/include/mlir/Tools/mlir-tblgen/MlirTblgenMain.h
@@ -0,0 +1,26 @@
+//===- MlirTblgenMain.h - MLIR Tablegen Driver main -------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Main entry function for mlir-tblgen for when built as standalone binary.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIR_TBLGEN_MLIRTBLGENMAIN_H
+#define MLIR_TOOLS_MLIR_TBLGEN_MLIRTBLGENMAIN_H
+
+namespace mlir {
+/// Main Program for tools like 'mlir-tblgen' with custom backends. To add
+/// a new backend, simply create a new 'mlir::GenRegistration' global variable.
+/// See its documentation for more info.
+///
+/// The 'argc' and 'argv' arguments are simply forwarded from a main function.
+/// The return value is the exit code from llvm::TableGenMain.
+int MlirTblgenMain(int argc, char **argv);
+} // namespace mlir
+
+#endif // MLIR_TOOLS_MLIR_TBLGEN_MLIRTBLGENMAIN_H
diff --git a/mlir/lib/TableGen/CMakeLists.txt b/mlir/lib/TableGen/CMakeLists.txt
index bb522d7d03f4..55a2e355d6a3 100644
--- a/mlir/lib/TableGen/CMakeLists.txt
+++ b/mlir/lib/TableGen/CMakeLists.txt
@@ -17,6 +17,7 @@ llvm_add_library(MLIRTableGen STATIC
   Constraint.cpp
   Dialect.cpp
   Format.cpp
+  GenInfo.cpp
   Interfaces.cpp
   Operator.cpp
   Pass.cpp
diff --git a/mlir/lib/TableGen/GenInfo.cpp b/mlir/lib/TableGen/GenInfo.cpp
new file mode 100644
index 000000000000..62a3d5283ef5
--- /dev/null
+++ b/mlir/lib/TableGen/GenInfo.cpp
@@ -0,0 +1,41 @@
+//===- GenInfo.cpp - Generator info -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/TableGen/GenInfo.h"
+
+#include "mlir/TableGen/GenNameParser.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ManagedStatic.h"
+
+using namespace mlir;
+
+static llvm::ManagedStatic<std::vector<GenInfo>> generatorRegistry;
+
+GenRegistration::GenRegistration(StringRef arg, StringRef description,
+                                 const GenFunction &function) {
+  generatorRegistry->emplace_back(arg, description, function);
+}
+
+GenNameParser::GenNameParser(llvm::cl::Option &opt)
+    : llvm::cl::parser<const GenInfo *>(opt) {
+  for (const auto &kv : *generatorRegistry) {
+    addLiteralOption(kv.getGenArgument(), &kv, kv.getGenDescription());
+  }
+}
+
+void GenNameParser::printOptionInfo(const llvm::cl::Option &o,
+                                    size_t globalWidth) const {
+  GenNameParser *tp = const_cast<GenNameParser *>(this);
+  llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),
+                       [](const GenNameParser::OptionInfo *vT1,
+                          const GenNameParser::OptionInfo *vT2) {
+                         return vT1->Name.compare(vT2->Name);
+                       });
+  using llvm::cl::parser;
+  parser<const GenInfo *>::printOptionInfo(o, globalWidth);
+}
diff --git a/mlir/lib/Tools/CMakeLists.txt b/mlir/lib/Tools/CMakeLists.txt
index 57e570437e8b..6dab371771a5 100644
--- a/mlir/lib/Tools/CMakeLists.txt
+++ b/mlir/lib/Tools/CMakeLists.txt
@@ -3,6 +3,7 @@ add_subdirectory(mlir-lsp-server)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
+add_subdirectory(mlir-tblgen)
 add_subdirectory(mlir-translate)
 add_subdirectory(PDLL)
 add_subdirectory(tblgen-lsp-server)
diff --git a/mlir/lib/Tools/mlir-tblgen/CMakeLists.txt b/mlir/lib/Tools/mlir-tblgen/CMakeLists.txt
new file mode 100644
index 000000000000..d1a83cf9fb10
--- /dev/null
+++ b/mlir/lib/Tools/mlir-tblgen/CMakeLists.txt
@@ -0,0 +1,18 @@
+llvm_add_library(MLIRTblgenLib STATIC
+  MlirTblgenMain.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-tblgen
+
+  DISABLE_LLVM_LINK_LLVM_DYLIB
+
+  LINK_COMPONENTS
+  TableGen
+
+  LINK_LIBS PUBLIC
+  MLIRTableGen
+  )
+
+mlir_check_all_link_libraries(MLIRTblgenLib)
+
+add_mlir_library_install(MLIRTblgenLib)
diff --git a/mlir/lib/Tools/mlir-tblgen/MlirTblgenMain.cpp b/mlir/lib/Tools/mlir-tblgen/MlirTblgenMain.cpp
new file mode 100644
index 000000000000..d5ed91e9b2b9
--- /dev/null
+++ b/mlir/lib/Tools/mlir-tblgen/MlirTblgenMain.cpp
@@ -0,0 +1,135 @@
+//===- MlirTblgenMain.cpp - MLIR Tablegen Driver main -----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Main entry function for mlir-tblgen for when built as standalone binary.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Tools/mlir-tblgen/MlirTblgenMain.h"
+
+#include "mlir/TableGen/GenInfo.h"
+#include "mlir/TableGen/GenNameParser.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/Signals.h"
+#include "llvm/TableGen/Error.h"
+#include "llvm/TableGen/Main.h"
+#include "llvm/TableGen/Record.h"
+
+using namespace mlir;
+using namespace llvm;
+
+enum DeprecatedAction { None, Warn, Error };
+
+static DeprecatedAction actionOnDeprecatedValue;
+
+// Returns if there is a use of `init` in `record`.
+static bool findUse(Record &record, Init *init,
+                    llvm::DenseMap<Record *, bool> &known) {
+  auto it = known.find(&record);
+  if (it != known.end())
+    return it->second;
+
+  auto memoize = [&](bool val) {
+    known[&record] = val;
+    return val;
+  };
+
+  for (const RecordVal &val : record.getValues()) {
+    Init *valInit = val.getValue();
+    if (valInit == init)
+      return true;
+    if (auto *di = dyn_cast<DefInit>(valInit)) {
+      if (findUse(*di->getDef(), init, known))
+        return memoize(true);
+    } else if (auto *di = dyn_cast<DagInit>(valInit)) {
+      for (Init *arg : di->getArgs())
+        if (auto *di = dyn_cast<DefInit>(arg))
+          if (findUse(*di->getDef(), init, known))
+            return memoize(true);
+    } else if (ListInit *li = dyn_cast<ListInit>(valInit)) {
+      for (Init *jt : li->getValues())
+        if (jt == init)
+          return memoize(true);
+    }
+  }
+  return memoize(false);
+}
+
+static void warnOfDeprecatedUses(RecordKeeper &records) {
+  // This performs a direct check for any def marked as deprecated and then
+  // finds all uses of deprecated def. Deprecated defs are not expected to be
+  // either numerous or long lived.
+  bool deprecatedDefsFounds = false;
+  for (auto &it : records.getDefs()) {
+    const RecordVal *r = it.second->getValue("odsDeprecated");
+    if (!r || !r->getValue())
+      continue;
+
+    llvm::DenseMap<Record *, bool> hasUse;
+    if (auto *si = dyn_cast<StringInit>(r->getValue())) {
+      for (auto &jt : records.getDefs()) {
+        // Skip anonymous defs.
+        if (jt.second->isAnonymous())
+          continue;
+        // Skip all outside main file to avoid flagging redundantly.
+        unsigned buf =
+            SrcMgr.FindBufferContainingLoc(jt.second->getLoc().front());
+        if (buf != SrcMgr.getMainFileID())
+          continue;
+
+        if (findUse(*jt.second, it.second->getDefInit(), hasUse)) {
+          PrintWarning(jt.second->getLoc(),
+                       "Using deprecated def `" + it.first + "`");
+          PrintNote(si->getAsUnquotedString());
+          deprecatedDefsFounds = true;
+        }
+      }
+    }
+  }
+  if (deprecatedDefsFounds &&
+      actionOnDeprecatedValue == DeprecatedAction::Error)
+    PrintFatalNote("Error'ing out due to deprecated defs");
+}
+
+// Generator to invoke.
+static const mlir::GenInfo *generator;
+
+// TableGenMain requires a function pointer so this function is passed in which
+// simply wraps the call to the generator.
+static bool mlirTableGenMain(raw_ostream &os, RecordKeeper &records) {
+  if (actionOnDeprecatedValue != DeprecatedAction::None)
+    warnOfDeprecatedUses(records);
+
+  if (!generator) {
+    os << records;
+    return false;
+  }
+  return generator->invoke(records, os);
+}
+
+int mlir::MlirTblgenMain(int argc, char **argv) {
+
+  llvm::InitLLVM y(argc, argv);
+
+  llvm::cl::opt<DeprecatedAction, true> actionOnDeprecated(
+      "on-deprecated", llvm::cl::desc("Action to perform on deprecated def"),
+      llvm::cl::values(
+          clEnumValN(DeprecatedAction::None, "none", "No action"),
+          clEnumValN(DeprecatedAction::Warn, "warn", "Warn on use"),
+          clEnumValN(DeprecatedAction::Error, "error", "Error on use")),
+      cl::location(actionOnDeprecatedValue), llvm::cl::init(Warn));
+
+  llvm::cl::opt<const mlir::GenInfo *, true, mlir::GenNameParser> generator(
+      "", llvm::cl::desc("Generator to run"), cl::location(::generator));
+
+  cl::ParseCommandLineOptions(argc, argv);
+
+  return TableGenMain(argv[0], &mlirTableGenMain);
+}
diff --git a/mlir/tools/mlir-tblgen/CMakeLists.txt b/mlir/tools/mlir-tblgen/CMakeLists.txt
index 43df44564e91..e3e3ed91515d 100644
--- a/mlir/tools/mlir-tblgen/CMakeLists.txt
+++ b/mlir/tools/mlir-tblgen/CMakeLists.txt
@@ -33,6 +33,6 @@ set_target_properties(mlir-tblgen PROPERTIES FOLDER "Tablegenning")
 target_link_libraries(mlir-tblgen
   PRIVATE
   MLIRSupportIndentedOstream
-  MLIRTableGen)
+  MLIRTblgenLib)
 
 mlir_check_all_link_libraries(mlir-tblgen)
diff --git a/mlir/tools/mlir-tblgen/mlir-tblgen.cpp b/mlir/tools/mlir-tblgen/mlir-tblgen.cpp
index cbc20c4a1af7..6c4b619598eb 100644
--- a/mlir/tools/mlir-tblgen/mlir-tblgen.cpp
+++ b/mlir/tools/mlir-tblgen/mlir-tblgen.cpp
@@ -11,56 +11,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/TableGen/GenInfo.h"
-#include "mlir/TableGen/GenNameParser.h"
-#include "llvm/ADT/StringExtras.h"
-#include "llvm/Support/CommandLine.h"
-#include "llvm/Support/FormatVariadic.h"
-#include "llvm/Support/InitLLVM.h"
-#include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/Signals.h"
-#include "llvm/TableGen/Error.h"
-#include "llvm/TableGen/Main.h"
+#include "mlir/Tools/mlir-tblgen/MlirTblgenMain.h"
 #include "llvm/TableGen/Record.h"
-#include "llvm/TableGen/TableGenBackend.h"
 
 using namespace llvm;
 using namespace mlir;
 
-enum DeprecatedAction { None, Warn, Error };
-llvm::cl::opt<DeprecatedAction> actionOnDeprecated(
-    "on-deprecated", llvm::cl::init(Warn),
-    llvm::cl::desc("Action to perform on deprecated def"),
-    llvm::cl::values(clEnumValN(DeprecatedAction::None, "none", "No action"),
-                     clEnumValN(DeprecatedAction::Warn, "warn", "Warn on use"),
-                     clEnumValN(DeprecatedAction::Error, "error",
-                                "Error on use")));
-
-static llvm::ManagedStatic<std::vector<GenInfo>> generatorRegistry;
-
-mlir::GenRegistration::GenRegistration(StringRef arg, StringRef description,
-                                       const GenFunction &function) {
-  generatorRegistry->emplace_back(arg, description, function);
-}
-
-GenNameParser::GenNameParser(llvm::cl::Option &opt)
-    : llvm::cl::parser<const GenInfo *>(opt) {
-  for (const auto &kv : *generatorRegistry) {
-    addLiteralOption(kv.getGenArgument(), &kv, kv.getGenDescription());
-  }
-}
-
-void GenNameParser::printOptionInfo(const llvm::cl::Option &o,
-                                    size_t globalWidth) const {
-  GenNameParser *tp = const_cast<GenNameParser *>(this);
-  llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),
-                       [](const GenNameParser::OptionInfo *vT1,
-                          const GenNameParser::OptionInfo *vT2) {
-                         return vT1->Name.compare(vT2->Name);
-                       });
-  using llvm::cl::parser;
-  parser<const GenInfo *>::printOptionInfo(o, globalWidth);
-}
-
 // Generator that prints records.
 GenRegistration printRecords("print-records", "Print all records to stdout",
                              [](const RecordKeeper &records, raw_ostream &os) {
@@ -68,96 +24,4 @@ GenRegistration printRecords("print-records", "Print all records to stdout",
                                return false;
                              });
 
-// Generator to invoke.
-const mlir::GenInfo *generator;
-
-// Returns if there is a use of `init` in `record`.
-bool findUse(Record &record, Init *init,
-             llvm::DenseMap<Record *, bool> &known) {
-  auto it = known.find(&record);
-  if (it != known.end())
-    return it->second;
-
-  auto memoize = [&](bool val) {
-    known[&record] = val;
-    return val;
-  };
-
-  for (const RecordVal &val : record.getValues()) {
-    Init *valInit = val.getValue();
-    if (valInit == init)
-      return true;
-    if (auto *di = dyn_cast<DefInit>(valInit)) {
-      if (findUse(*di->getDef(), init, known))
-        return memoize(true);
-    } else if (auto *di = dyn_cast<DagInit>(valInit)) {
-      for (Init *arg : di->getArgs())
-        if (auto *di = dyn_cast<DefInit>(arg))
-          if (findUse(*di->getDef(), init, known))
-            return memoize(true);
-    } else if (ListInit *li = dyn_cast<ListInit>(valInit)) {
-      for (Init *jt : li->getValues())
-        if (jt == init)
-          return memoize(true);
-    }
-  }
-  return memoize(false);
-}
-
-void warnOfDeprecatedUses(RecordKeeper &records) {
-  // This performs a direct check for any def marked as deprecated and then
-  // finds all uses of deprecated def. Deprecated defs are not expected to be
-  // either numerous or long lived.
-  bool deprecatedDefsFounds = false;
-  for (auto &it : records.getDefs()) {
-    const RecordVal *r = it.second->getValue("odsDeprecated");
-    if (!r || !r->getValue())
-      continue;
-
-    llvm::DenseMap<Record *, bool> hasUse;
-    if (auto *si = dyn_cast<StringInit>(r->getValue())) {
-      for (auto &jt : records.getDefs()) {
-        // Skip anonymous defs.
-        if (jt.second->isAnonymous())
-          continue;
-        // Skip all outside main file to avoid flagging redundantly.
-        unsigned buf =
-            SrcMgr.FindBufferContainingLoc(jt.second->getLoc().front());
-        if (buf != SrcMgr.getMainFileID())
-          continue;
-
-        if (findUse(*jt.second, it.second->getDefInit(), hasUse)) {
-          PrintWarning(jt.second->getLoc(),
-                       "Using deprecated def `" + it.first + "`");
-          PrintNote(si->getAsUnquotedString());
-          deprecatedDefsFounds = true;
-        }
-      }
-    }
-  }
-  if (deprecatedDefsFounds && actionOnDeprecated == DeprecatedAction::Error)
-    PrintFatalNote("Error'ing out due to deprecated defs");
-}
-
-// TableGenMain requires a function pointer so this function is passed in which
-// simply wraps the call to the generator.
-static bool mlirTableGenMain(raw_ostream &os, RecordKeeper &records) {
-  if (actionOnDeprecated != DeprecatedAction::None)
-    warnOfDeprecatedUses(records);
-
-  if (!generator) {
-    os << records;
-    return false;
-  }
-  return generator->invoke(records, os);
-}
-
-int main(int argc, char **argv) {
-  llvm::InitLLVM y(argc, argv);
-  llvm::cl::opt<const mlir::GenInfo *, false, mlir::GenNameParser> generator(
-      "", llvm::cl::desc("Generator to run"));
-  cl::ParseCommandLineOptions(argc, argv);
-  ::generator = generator.getValue();
-
-  return TableGenMain(argv[0], &mlirTableGenMain);
-}
+int main(int argc, char **argv) { return MlirTblgenMain(argc, argv); }