Support file-to-file translation in mlir-translate

Existing translations are either from MLIR or to MLIR. To support
cases like round-tripping some external format via MLIR, one must
chain two mlir-translate invocations together using pipes. This
can be problematic to support -split-input-file in mlir-translate
given that it won't work across pipes.

Motivated by the above, this CL adds another translation category
that allows file to file. This gives users more freedom.

PiperOrigin-RevId: 269636438

mlir/include/mlir/Translation.h

@@ -50,12 +50,14 @@ using TranslateFromMLIRFunction =
 
 /// Interface of the function that performs file-to-file translation involving
 /// MLIR. The input file is held in the given MemoryBuffer; the output file
-/// should be written to the given raw_ostream.
+/// should be written to the given raw_ostream. The implementation should create
+/// all MLIR constructs needed during the process inside the given context. This
+/// can be used for round-tripping external formats through the MLIR system.
 using TranslateFunction =
     std::function<LogicalResult(std::unique_ptr<llvm::MemoryBuffer> input,
                                 llvm::raw_ostream &output, MLIRContext *)>;
 
-/// Use Translate[To|From]MLIRRegistration as a global initialiser that
+/// Use Translate[ToMLIR|FromMLIR|]Registration as a global initialiser that
 /// registers a function and associates it with name. This requires that a
 /// translation has not been registered to a given name.
 ///
@@ -74,12 +76,17 @@ struct TranslateFromMLIRRegistration {
   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<TranslateToMLIRFunction> &getTranslationToMLIRRegistry();
 const llvm::StringMap<TranslateFromMLIRFunction> &
 getTranslationFromMLIRRegistry();
+const llvm::StringMap<TranslateFunction> &getTranslationRegistry();
 
 } // namespace mlir
 

mlir/lib/Support/TranslateClParser.cpp

@@ -48,10 +48,12 @@ TranslationParser::TranslationParser(llvm::cl::Option &opt)
     : llvm::cl::parser<const TranslateFunction *>(opt) {
   const auto &toMLIRRegistry = getTranslationToMLIRRegistry();
   const auto &fromMLIRRegistry = getTranslationFromMLIRRegistry();
+  const auto &fileToFileRegistry = getTranslationRegistry();
 
   // Reserve the required capacity upfront so that pointers are not
   // invalidated on reallocation.
-  wrapperStorage.reserve(toMLIRRegistry.size() + fromMLIRRegistry.size());
+  wrapperStorage.reserve(toMLIRRegistry.size() + fromMLIRRegistry.size() +
+                         fileToFileRegistry.size());
   for (const auto &kv : toMLIRRegistry) {
     TranslateToMLIRFunction function = kv.second;
     TranslateFunction wrapper =
@@ -85,6 +87,10 @@ TranslationParser::TranslationParser(llvm::cl::Option &opt)
 
     addLiteralOption(kv.first(), &wrapperStorage.back(), kv.first());
   }
+  for (const auto &kv : fileToFileRegistry) {
+    wrapperStorage.emplace_back(kv.second);
+    addLiteralOption(kv.first(), &wrapperStorage.back(), kv.first());
+  }
 }
 
 void TranslationParser::printOptionInfo(const llvm::cl::Option &O,

mlir/lib/Translation/Translation.cpp

@@ -42,6 +42,13 @@ getMutableTranslationFromMLIRRegistry() {
   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() {
+  static llvm::StringMap<TranslateFunction> translationRegistry;
+  return translationRegistry;
+}
+
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, const TranslateToMLIRFunction &function) {
   auto &translationToMLIRRegistry = getMutableTranslationToMLIRRegistry();
@@ -64,6 +71,17 @@ TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
   translationFromMLIRRegistry[name] = function;
 }
 
+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<TranslateToMLIRFunction> &
@@ -75,3 +93,7 @@ const llvm::StringMap<TranslateFromMLIRFunction> &
 mlir::getTranslationFromMLIRRegistry() {
   return getMutableTranslationFromMLIRRegistry();
 }
+
+const llvm::StringMap<TranslateFunction> &mlir::getTranslationRegistry() {
+  return getMutableTranslationRegistry();
+}