[mlir] Refactor and cleanup the translation facilities.

[rriddle]

Summary

This revision performs several cleanups on the translation infra:

• Removes the TranslateCLParser library and consolidates into Translation
  • This was a weird library that existed in Support, and didn't really justify being a standalone library.
• Cleans up the internal registration and consolidates all of the translation functions within one registry.

[mehdi_amini]

What happens with TranslateClParser.cpp? It isn't showed as deleted in this revision it seems.

[mehdi_amini]

mlir/lib/Translation/Translation.cpp

62

It'd be nice to have an option to disable the verifier when testing.

[rriddle]

Remove files

mlir/include/mlir/Translation.h

 //===- Translation.h - Translation registry ---------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Registry for user-provided translations.
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_TRANSLATION_H
 #define MLIR_TRANSLATION_H
 
-#include "llvm/ADT/StringMap.h"
-
-#include <memory>
+#include "llvm/Support/CommandLine.h"
 
 namespace llvm {
 class MemoryBuffer;
 class SourceMgr;
 class StringRef;
 } // namespace llvm
 
 namespace mlir {
TranslateFromMLIRRegistration(llvm::StringRef name,
                                 const TranslateFromMLIRFunction &function);
 };
 struct TranslateRegistration {
   TranslateRegistration(llvm::StringRef name,
                         const TranslateFunction &function);
 };
 /// \}
 
-/// Get a read-only reference to the translator registry.
-const llvm::StringMap<TranslateSourceMgrToMLIRFunction> &
-getTranslationToMLIRRegistry();
-const llvm::StringMap<TranslateFromMLIRFunction> &
-getTranslationFromMLIRRegistry();
-const llvm::StringMap<TranslateFunction> &getTranslationRegistry();
+/// A command line parser for translation functions.
+struct TranslationParser : public llvm::cl::parser<const TranslateFunction *> {
+  TranslationParser(llvm::cl::Option &opt);
 
+  void printOptionInfo(const llvm::cl::Option &o,
+                       size_t globalWidth) const override;
+};
 } // namespace mlir
 
 #endif // MLIR_TRANSLATION_H

mlir/lib/Support/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   FileUtilities.cpp
   JitRunner.cpp
   MlirOptMain.cpp
   StorageUniquer.cpp
   ToolUtilities.cpp
-  TranslateClParser.cpp
 )
 
 add_mlir_library(MLIRSupport
   FileUtilities.cpp
   StorageUniquer.cpp
   ToolUtilities.cpp
 
   ADDITIONAL_HEADER_DIRS

 target_link_libraries(MLIROptLib
   PUBLIC
   MLIRPass
   MLIRParser
   LLVMSupport
   MLIRSupport
   )
 
-add_mlir_library(MLIRTranslateClParser
-  TranslateClParser.cpp
-
-  ADDITIONAL_HEADER_DIRS
-  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
-  )
-target_link_libraries(MLIRTranslateClParser
-  PUBLIC
-  LLVMSupport
-  MLIRIR
-  MLIRTranslation
-  MLIRParser)
-
 add_llvm_library(MLIRJitRunner
   JitRunner.cpp
 )
 target_link_libraries(MLIRJitRunner
   PUBLIC
   MLIRExecutionEngine
   MLIRIR
   MLIRParser
   MLIRStandardOps
   MLIRTargetLLVMIR
   MLIRTransforms
   MLIRStandardToLLVM
   MLIRSupport
   LLVMCore
   LLVMSupport
 )

mlir/lib/Translation/CMakeLists.txt

 add_mlir_library(MLIRTranslation
   Translation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Translation
   )
 target_link_libraries(MLIRTranslation
   PUBLIC
   LLVMSupport
   MLIRIR
+  MLIRParser
   )

mlir/lib/Translation/Translation.cpp

 //===- Translation.cpp - Translation registry -----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Definitions of the translation registry.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Translation.h"
+#include "mlir/Analysis/Verifier.h"
 #include "mlir/IR/Module.h"
+#include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
-#include "mlir/Support/LogicalResult.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
-// Get the mutable static map between registered "to MLIR" translations and the
-// TranslateToMLIRFunctions that perform those translations.
-static llvm::StringMap<TranslateSourceMgrToMLIRFunction> &
-getMutableTranslationToMLIRRegistry() {
-  static llvm::StringMap<TranslateSourceMgrToMLIRFunction>
-      translationToMLIRRegistry;
-  return translationToMLIRRegistry;
-}
-// Get the mutable static map between registered "from MLIR" translations and
-// the TranslateFromMLIRFunctions that perform those translations.
-static llvm::StringMap<TranslateFromMLIRFunction> &
-getMutableTranslationFromMLIRRegistry() {
-  static llvm::StringMap<TranslateFromMLIRFunction> translationFromMLIRRegistry;
-  return translationFromMLIRRegistry;
-}
-
-// Get the mutable static map between registered file-to-file MLIR translations
-// and the TranslateFunctions that perform those translations.
-static llvm::StringMap<TranslateFunction> &getMutableTranslationRegistry() {
+//===----------------------------------------------------------------------===//
+// Translation Registry
+//===----------------------------------------------------------------------===//
+
+/// Get the mutable static map between registered file-to-file MLIR translations
+/// and the TranslateFunctions that perform those translations.
+static llvm::StringMap<TranslateFunction> &getTranslationRegistry() {
   static llvm::StringMap<TranslateFunction> translationRegistry;
   return translationRegistry;
 }
 
+/// Register the given translation.
+static void registerTranslation(StringRef name,
+                                const TranslateFunction &function) {
+  auto &translationRegistry = getTranslationRegistry();
+  if (translationRegistry.find(name) != translationRegistry.end())
+    llvm::report_fatal_error(
+        "Attempting to overwrite an existing <file-to-file> function");
+  assert(function &&
+         "Attempting to register an empty translate <file-to-file> function");
+  translationRegistry[name] = function;
+}
+
+TranslateRegistration::TranslateRegistration(
+    StringRef name, const TranslateFunction &function) {
+  registerTranslation(name, function);
+}
+
+//===----------------------------------------------------------------------===//
+// Translation to MLIR
+//===----------------------------------------------------------------------===//
+
 // Puts `function` into the to-MLIR translation registry unless there is already
 // a function registered for the same name.
 static void registerTranslateToMLIRFunction(
     StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
-  auto &translationToMLIRRegistry = getMutableTranslationToMLIRRegistry();
-  if (translationToMLIRRegistry.find(name) != translationToMLIRRegistry.end())
-    llvm::report_fatal_error(
-        "Attempting to overwrite an existing <to> function");
-  assert(function && "Attempting to register an empty translate <to> function");
-  translationToMLIRRegistry[name] = function;
+  auto wrappedFn = [function](llvm::SourceMgr &sourceMgr, raw_ostream &output,
+                              MLIRContext *context) {
+    OwningModuleRef module = function(sourceMgr, context);
+    if (!module || failed(verify(*module)))
+      return failure();
+    module->print(output);
+    return success();
+  };
+  registerTranslation(name, wrappedFn);
 }
 
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
   registerTranslateToMLIRFunction(name, function);
 }
 
-// Wraps `function` with a lambda that extracts a StringRef from a source
-// manager and registers the wrapper lambda as a to-MLIR conversion.
+/// Wraps `function` with a lambda that extracts a StringRef from a source
+/// manager and registers the wrapper lambda as a to-MLIR conversion.
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, const TranslateStringRefToMLIRFunction &function) {
-  auto translationFunction = [function](llvm::SourceMgr &sourceMgr,
-                                        MLIRContext *ctx) {
+  registerTranslateToMLIRFunction(
+      name, [function](llvm::SourceMgr &sourceMgr, MLIRContext *ctx) {
         const llvm::MemoryBuffer *buffer =
             sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
         return function(buffer->getBuffer(), ctx);
-  };
-  registerTranslateToMLIRFunction(name, translationFunction);
+      });
 }
 
+//===----------------------------------------------------------------------===//
+// Translation from MLIR
+//===----------------------------------------------------------------------===//
+
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, const TranslateFromMLIRFunction &function) {
-  auto &translationFromMLIRRegistry = getMutableTranslationFromMLIRRegistry();
-  if (translationFromMLIRRegistry.find(name) !=
-      translationFromMLIRRegistry.end())
-    llvm::report_fatal_error(
-        "Attempting to overwrite an existing <from> function");
-  assert(function &&
-         "Attempting to register an empty translate <from> function");
-  translationFromMLIRRegistry[name] = function;
+  registerTranslation(name, [function](llvm::SourceMgr &sourceMgr,
+                                       raw_ostream &output,
+                                       MLIRContext *context) {
+    auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
+    if (!module)
+      return failure();
+    return function(module.get(), output);
+  });
 }
 
-TranslateRegistration::TranslateRegistration(
-    StringRef name, const TranslateFunction &function) {
-  auto &translationRegistry = getMutableTranslationRegistry();
-  if (translationRegistry.find(name) != translationRegistry.end())
-    llvm::report_fatal_error(
-        "Attempting to overwrite an existing <file-to-file> function");
-  assert(function &&
-         "Attempting to register an empty translate <file-to-file> function");
-  translationRegistry[name] = function;
-}
-
-// Merely add the const qualifier to the mutable registry so that external users
-// cannot modify it.
-const llvm::StringMap<TranslateSourceMgrToMLIRFunction> &
-mlir::getTranslationToMLIRRegistry() {
-  return getMutableTranslationToMLIRRegistry();
-}
+//===----------------------------------------------------------------------===//
+// Translation Parser
+//===----------------------------------------------------------------------===//
 
-const llvm::StringMap<TranslateFromMLIRFunction> &
-mlir::getTranslationFromMLIRRegistry() {
-  return getMutableTranslationFromMLIRRegistry();
+TranslationParser::TranslationParser(llvm::cl::Option &opt)
+    : llvm::cl::parser<const TranslateFunction *>(opt) {
+  for (const auto &kv : getTranslationRegistry())
+    addLiteralOption(kv.first(), &kv.second, kv.first());
 }
 
-const llvm::StringMap<TranslateFunction> &mlir::getTranslationRegistry() {
-  return getMutableTranslationRegistry();
+void TranslationParser::printOptionInfo(const llvm::cl::Option &o,
+                                        size_t globalWidth) const {
+  TranslationParser *tp = const_cast<TranslationParser *>(this);
+  llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),
+                       [](const TranslationParser::OptionInfo *lhs,
+                          const TranslationParser::OptionInfo *rhs) {
+                         return lhs->Name.compare(rhs->Name);
+                       });
+  llvm::cl::parser<const TranslateFunction *>::printOptionInfo(o, globalWidth);
 }

mlir/tools/mlir-translate/CMakeLists.txt

set(FULL_LIBS
   MLIRTargetNVVMIR
   MLIRTargetROCDLIR
 )
 add_llvm_tool(mlir-translate
   mlir-translate.cpp
 )
 llvm_update_compile_flags(mlir-translate)
 whole_archive_link(mlir-translate ${FULL_LIBS})
-target_link_libraries(mlir-translate PRIVATE MLIRIR MLIRTranslateClParser ${LIBS} LLVMSupport)
+target_link_libraries(mlir-translate PRIVATE MLIRIR MLIRTranslation ${LIBS} LLVMSupport)

mlir/tools/mlir-translate/mlir-translate.cpp

 //===- mlir-translate.cpp - MLIR Translate Driver -------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This is a command line utility that translates a file from/to MLIR using one
 // of the registered translations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/ToolUtilities.h"
-#include "mlir/Support/TranslateClParser.h"
+#include "mlir/Translation.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
 static llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional,