Add an optional dialect_registration callback for "translations" registered with mlir-translate

[mehdi_amini]

Summary

This will allow out-of-tree translation to register the dialects they expect
to see in their input, on the model of getDependentDialects() for passes.

[mehdi_amini]

Minor tweak

[rriddle]

mlir/lib/Dialect/SPIRV/Serialization/TranslateRegistration.cpp

127

Why is this needed?

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

117

Did you ever add the variadic template variant?

[mehdi_amini]

address comments

[mehdi_amini]

mlir/lib/Dialect/SPIRV/Serialization/TranslateRegistration.cpp

127

It is registered as a TranslateRegistration and not a TranslateFromMLIRRegistration

The contract is different as this function is doing the parsing itself instead of getting a Module as input.

I just pushed https://github.com/llvm/llvm-project/commit/12541b5ed59d and I'll update this patch to align this translation on the others.

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

117

Yes, updated here (and everywhere)

mlir/include/mlir/Translation.h

 struct TranslateToMLIRRegistration {
   TranslateToMLIRRegistration(llvm::StringRef name,
                               const TranslateSourceMgrToMLIRFunction &function);
   TranslateToMLIRRegistration(llvm::StringRef name,
                               const TranslateStringRefToMLIRFunction &function);
 };
 
 struct TranslateFromMLIRRegistration {
-  TranslateFromMLIRRegistration(llvm::StringRef name,
-                                const TranslateFromMLIRFunction &function);
+  TranslateFromMLIRRegistration(
+      llvm::StringRef name, const TranslateFromMLIRFunction &function,
+      std::function<void(DialectRegistry &)> dialectRegistration =
+          [](DialectRegistry &) {});
 };
 struct TranslateRegistration {
   TranslateRegistration(llvm::StringRef name,
                         const TranslateFunction &function);
 };
 /// \}
 
 /// A command line parser for translation functions.

mlir/lib/Dialect/SPIRV/Serialization/TranslateRegistration.cpp

 //===- TranslateRegistration.cpp - hooks to mlir-translate ----------------===//
 //
 // 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 file implements a translation from SPIR-V binary module to MLIR SPIR-V
 // ModuleOp.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Translation.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SMLoc.h"
output.write(reinterpret_cast<char *>(binary.data()),
                binary.size() * sizeof(uint32_t));
 
   return mlir::success();
 }
 
 namespace mlir {
 void registerToSPIRVTranslation() {
   TranslateFromMLIRRegistration toBinary(
-      "serialize-spirv", [](ModuleOp module, raw_ostream &output) {
+      "serialize-spirv",
+      [](ModuleOp module, raw_ostream &output) {
         return serializeModule(module, output);
+      },
+      [](DialectRegistry &registry) {
+        registry.insert<spirv::SPIRVDialect>();
       });
 }
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Round-trip registration
 //===----------------------------------------------------------------------===//
 
 static LogicalResult roundTripModule(llvm::SourceMgr &sourceMgr,
                                      bool emitDebugInfo, raw_ostream &output,
                                      MLIRContext *context) {
+  context->getDialectRegistry().insert<spirv::SPIRVDialect>();
   // Parse an MLIR module from the source manager.
   auto srcModule = OwningModuleRef(parseSourceFile(sourceMgr, context));
   if (!srcModule)
     return failure();
 
   SmallVector<uint32_t, 0> binary;
 
   auto spirvModules = srcModule->getOps<spirv::ModuleOp>();

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 mlir::translateModuleToLLVMIR(ModuleOp m, llvm::LLVMContext &llvmContext,
                               StringRef name) {
   return LLVM::ModuleTranslation::translateModule<>(m, llvmContext, name);
 }
 
 namespace mlir {
 void registerToLLVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
-      "mlir-to-llvmir", [](ModuleOp module, raw_ostream &output) {
+      "mlir-to-llvmir",
+      [](ModuleOp module, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = LLVM::ModuleTranslation::translateModule<>(
             module, llvmContext, "LLVMDialectModule");
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
-      });
+      },
+      [](DialectRegistry &registry) { registry.insert<LLVM::LLVMDialect>(); });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp

mlir::translateModuleToNVVMIR(Operation *m, llvm::LLVMContext &llvmContext,
   }
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToNVVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
-      "mlir-to-nvvmir", [](ModuleOp module, raw_ostream &output) {
+      "mlir-to-nvvmir",
+      [](ModuleOp module, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = mlir::translateModuleToNVVMIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
+      },
+      [](DialectRegistry &registry) {
+        registry.insert<NVVM::NVVMDialect>();
+        registry.insert<LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

mlir::translateModuleToROCDLIR(Operation *m, llvm::LLVMContext &llvmContext,
   }
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToROCDLIRTranslation() {
   TranslateFromMLIRRegistration registration(
-      "mlir-to-rocdlir", [](ModuleOp module, raw_ostream &output) {
+      "mlir-to-rocdlir",
+      [](ModuleOp module, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = mlir::translateModuleToROCDLIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
+      },
+      [](DialectRegistry &registry) {
+        registry.insert<ROCDL::ROCDLDialect>();
+        registry.insert<LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp

translateLLVMAVX512ModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
   return LLVM::ModuleTranslation::translateModule<LLVMAVX512ModuleTranslation>(
       m, llvmContext, name);
 }
 } // end namespace
 
 namespace mlir {
 void registerAVX512ToLLVMIRTranslation() {
   TranslateFromMLIRRegistration reg(
-      "avx512-mlir-to-llvmir", [](ModuleOp module, raw_ostream &output) {
+      "avx512-mlir-to-llvmir",
+      [](ModuleOp module, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = translateLLVMAVX512ModuleToLLVMIR(
             module, llvmContext, "LLVMDialectModule");
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
+      },
+      [](DialectRegistry &registry) {
+        registry.insert<LLVM::LLVMAVX512Dialect>();
+        registry.insert<LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Translation/Translation.cpp

registerTranslateToMLIRFunction(
       });
 }
 
 //===----------------------------------------------------------------------===//
 // Translation from MLIR
 //===----------------------------------------------------------------------===//
 
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
-    StringRef name, const TranslateFromMLIRFunction &function) {
-  registerTranslation(name, [function](llvm::SourceMgr &sourceMgr,
-                                       raw_ostream &output,
-                                       MLIRContext *context) {
+    StringRef name, const TranslateFromMLIRFunction &function,
+    std::function<void(DialectRegistry &)> dialectRegistration) {
+  registerTranslation(name, [function, dialectRegistration](
+                                llvm::SourceMgr &sourceMgr, raw_ostream &output,
+    dialectRegistration(context->getDialectRegistry());
     auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
     if (!module)
       return failure();
     return function(module.get(), output);
   });
 }
 
 //===----------------------------------------------------------------------===//
LogicalResult mlir::MlirTranslateMain(int argc, char **argv,
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   // Processes the memory buffer with a new MLIRContext.
   auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
                            raw_ostream &os) {
-    MLIRContext context;
+    MLIRContext context(false);
     context.printOpOnDiagnostic(!verifyDiagnostics);
     llvm::SourceMgr sourceMgr;
     sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), llvm::SMLoc());
 
     if (!verifyDiagnostics) {
       SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
       return (*translationRequested)(sourceMgr, os, &context);
     }

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

 static void registerTestTranslations() {
   registerTestRoundtripSPIRV();
   registerTestRoundtripDebugSPIRV();
 }
 
 int main(int argc, char **argv) {
   registerAllTranslations();
   registerTestTranslations();
-  // TODO: remove the global dialect registry
-  registerAllDialects();
   return failed(MlirTranslateMain(argc, argv, "MLIR Translation Testing Tool"));
 }