[mlir] Add translation parameter to mlir-cpu-runner.

[springerm]

Summary

The current mlir-cpu-runner implementation always uses
translateModuleToLLVMIR. With this change, the translation can be specified
from via a command line parameter. That makes it possible to write integration
tests for other dialects such as LLVMAVX512_Dialect.

Depends On D95678.

[rriddle]

I'm not sure I agree with the direction here. I don't think we should be hardcoding anything LLVM related in the Translation library.

[springerm]

I see. I wanted to reuse some of the existing TranslationRegistry infrastructure, but I can it change so that the entire new logic is in mlir/lib/ExecutionEngine.

[nicolasvasilache]

In D95732#2532068, @rriddle wrote:

I'm not sure I agree with the direction here. I don't think we should be hardcoding anything LLVM related in the Translation library.

Atm I am not aware of MLIR translating to anything else that LLVM, is your thinking that the Translation library has/will be the supporting library for translating from MLIR to future targets?
In that case, should we start a new binary mlir-translate-to-llvmir that is LLVM-specific?
Or do you have something else in mind?

[mehdi_amini]

In D95732#2532163, @nicolasvasilache wrote:

In D95732#2532068, @rriddle wrote:

I'm not sure I agree with the direction here. I don't think we should be hardcoding anything LLVM related in the Translation library.

I agree, this patch is really not clear to me right now.
I'd rather revisit our general translation mechanism to LLVM and the specificity of mlir-cpu-runner than changing the fundamental of mlir-translate.

Atm I am not aware of MLIR translating to anything else that LLVM, is your thinking that the Translation library has/will be the supporting library for translating from MLIR to future targets?

In tree we have at least SPIRV on top of my head. Out-of-tree it is used for emitting graphdef, TFLite flat buffer, HLO proto, etc.

mlir/include/mlir/Target/LLVMIR.h

38

Why is this changing? (the doc still says "module")

[springerm]

How about changing the commit as follows:

• In ConvertToLLVMIR.cpp / LLVMAVX512Intr.cpp, etc.: In addition to doing the TranslateFromMLIRRegistration, perform a second TranslateFromMLIRToLLVMRegistration that registers the corresponding translateModuleToLLVMIR / translateLLVMAVX512ModuleToLLVMIR function.
• TranslateFromMLIRToLLVMRegistration is defined in lib/ExecutionEngine/Translation.cpp, so there are no changes to mlir-translate.
• Define a new JitRunnerTranslationParser in lib/ExecutionEngine/Translation.cpp. Basically a copy of the TranslationParser of mlir-translate, but this one will be used only for mlir-cpu-runner.

This would avoid mixing mlir-translate with LLVM-specific stuff.

But it also duplicates the entire TranslationParser / Translate [...] Registration logic and translation passes in Target/LLVMIR/*.cpp will essentially be registered twice. Even though most TranslateFromMLIRRegistrations (apart from the SPIRV one) are just a wrapper around translate*ToLLVMIR, e.g.:

void registerAVX512ToLLVMIRTranslation() {
  TranslateFromMLIRRegistration reg(
      "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, LLVM::LLVMDialect>();
      });
}

Here, registerAVX512ToLLVMIRTranslation is just a wrapper around translateLLVMAVX512ModuleToLLVMIR (which calls LLVM::ModuleTranslation::translateModule). Actually, all of those functions in lib/Target/LLVMIR/*.cpp look the same, so there's a chance to remove some code duplication. (Which I attempted to do in TranslateFromMLIRToLLVMRegistration::TranslateFromMLIRToLLVMRegistration in this commit.)

Another approach would be to move the entire TranslationParser / Translate [...] Registration logic into a new component, which can be used by both mlir-cpu-runner and mlir-translate. Since they both perform "MLIR --> something" translations, it would make sense for them to share some code (such as the registration of available translations).

Any other ideas? What do you think?

[springerm]

mlir/include/mlir/Target/LLVMIR.h

38

Forgot to update the comment.

This function is inconsistent with the other translate*ToLLVMIR functions such as translateModuleToNVVMIR or translateLLVMArmSVEModuleToLLVMIR. These functions are all taking an Operation*.

By changing ModuleOP to Operation*, it was easier to register these functions with TranslationRegistry<TranslateMLIRToLLVMFunction> (now that they all have the same signature).

[nicolasvasilache]

mlir/include/mlir/Target/LLVMIR.h

38

This seems more like an oversight on the the operation that are called translateXXXModuleYYY and take an operation ?
I'd go the other way around: make all of those take a ModuleOp.

[mehdi_amini]

The root of the problem is that we have multiple dialects that are involved in a direct translation to LLVM IR without going to the LLVM dialect itself. This is a departure from the original design where everything should be transformed into the LLVM dialect before translation (which ensures also proper round-tripping from LLVM IR to LLVM dialect).

Assuming we want to endorse this new paradigm, I'd look into making the general Translate MLIR To LLVM pluggable to dialects that provide direct conversion.
To do this, I'd keep a single entry point mlir-to-llvmir and remove all the others (avx512-mlir-to-llvmir should not exist). The implementation LLVM::ModuleTranslation::convertOperation could be final. The implementation can simply lookup the dialect for each operation and dispatch the conversion through a new interface.

This way we don't rely on any global static registration (other than the existing MLIRContext) for the translation, new Dialect can easily plug into the LLVM translation, including out-of-tree dialects.

[ftynse]

In D95732#2533185, @mehdi_amini wrote:

The root of the problem is that we have multiple dialects that are involved in a direct translation to LLVM IR without going to the LLVM dialect itself. This is a departure from the original design where everything should be transformed into the LLVM dialect before translation (which ensures also proper round-tripping from LLVM IR to LLVM dialect).

Assuming we want to endorse this new paradigm, I'd look into making the general Translate MLIR To LLVM pluggable to dialects that provide direct conversion.
To do this, I'd keep a single entry point mlir-to-llvmir and remove all the others (avx512-mlir-to-llvmir should not exist). The implementation LLVM::ModuleTranslation::convertOperation could be final. The implementation can simply lookup the dialect for each operation and dispatch the conversion through a new interface.

This way we don't rely on any global static registration (other than the existing MLIRContext) for the translation, new Dialect can easily plug into the LLVM translation, including out-of-tree dialects.

I agree with this assessment. FWIW, we have de facto diverged from the original design some time ago. Specifically, when we decided to put all target-specific LLVM IR intrinsics into separate dialects (nvvm was actually the first, but avx, neon and sve are in the same situation). The reason for that is to avoid having thousands of unused ops in the main LLVM dialect with all the increased resource usage it entails. Then there is the OpenMP dialect, which uses OpenMPIRBuilder to emit loads of LLVM IR from a single MLIR operation bypassing the LLVM dialect. Furthermore, one may have OpenMP and, e.g., AVX simultaneously, so having -translate-something-to-llvm doesn't scale well because "something" may be a combination. I have been thinking about this in the background and my general inclination is to go towards more pluggability, similarly to what you suggest. Yes, making it easier to bypass the LLVM dialect when targeting LLVM IR may create wrong incentives, but it isn't that hard to create a new translation today.

There is one layering issue: if we want each dialect to know how to lower itself to LLVM IR, this creates a dependency from that dialect to LLVM IR core libraries and, potentially, to Translation. Arguably, this isn't something we want if we are just using the dialects, e.g., in mlir-opt. It took me significant time to get rid of the dependency on thread-unsafe LLVMContext in the LLVM dialect so I'd prefer not running a risk of putting that back in. I would be happy with some more dynamic approach.

[ftynse]

mlir/include/mlir/Target/LLVMIR.h

38

They can't. NVVM takes a gpu.module IIRC.

[nicolasvasilache]

In D95732#2533185, @mehdi_amini wrote:

The root of the problem is that we have multiple dialects that are involved in a direct translation to LLVM IR without going to the LLVM dialect itself. This is a departure from the original design where everything should be transformed into the LLVM dialect before translation (which ensures also proper round-tripping from LLVM IR to LLVM dialect).

There is one layering issue: if we want each dialect to know how to lower itself to LLVM IR, this creates a dependency from that dialect to LLVM IR core libraries and, potentially, to Translation. Arguably, this isn't something we want if we are just using the dialects, e.g., in mlir-opt. It took me significant time to get rid of the dependency on thread-unsafe LLVMContext in the LLVM dialect so I'd prefer not running a risk of putting that back in. I would be happy with some more dynamic approach.

An additional layering to consider here is that some intrinsics (e.g. ARM) but not all (e.g. AVX512) are resolved in an overloaded fashion by the LLVM Builder API Intrinsic::getDeclaration.
I.e. one can't reliably construct the "intrinsic function name" to build a llvm.call without either calling the LLVM Builder API or copy-pasting the internals, both of which seem unacceptable.

[ftynse]

In D95732#2533477, @nicolasvasilache wrote:

In D95732#2533185, @mehdi_amini wrote:

The root of the problem is that we have multiple dialects that are involved in a direct translation to LLVM IR without going to the LLVM dialect itself. This is a departure from the original design where everything should be transformed into the LLVM dialect before translation (which ensures also proper round-tripping from LLVM IR to LLVM dialect).

There is one layering issue: if we want each dialect to know how to lower itself to LLVM IR, this creates a dependency from that dialect to LLVM IR core libraries and, potentially, to Translation. Arguably, this isn't something we want if we are just using the dialects, e.g., in mlir-opt. It took me significant time to get rid of the dependency on thread-unsafe LLVMContext in the LLVM dialect so I'd prefer not running a risk of putting that back in. I would be happy with some more dynamic approach.

An additional layering to consider here is that some intrinsics (e.g. ARM) but not all (e.g. AVX512) are resolved in an overloaded fashion by the LLVM Builder API Intrinsic::getDeclaration.
I.e. one can't reliably construct the "intrinsic function name" to build a llvm.call without either calling the LLVM Builder API or copy-pasting the internals, both of which seem unacceptable.

I don't think we want to lower ops that correspond to intrinsics to LLVM dialect function calls. So this is the same layering issue of the "intrinsic-derived" dialect having a dependency on LLVM libraries.

[mehdi_amini]

In D95732#2533368, @ftynse wrote:

There is one layering issue: if we want each dialect to know how to lower itself to LLVM IR, this creates a dependency from that dialect to LLVM IR core libraries and, potentially, to Translation. Arguably, this isn't something we want if we are just using the dialects, e.g., in mlir-opt. It took me significant time to get rid of the dependency on thread-unsafe LLVMContext in the LLVM dialect so I'd prefer not running a risk of putting that back in. I would be happy with some more dynamic approach.

Ah this is an excellent point as well. This makes it tricky to rely on dialect interfaces, or we'd need maybe some injection into the Dialect registration for these dialects.

[springerm]

Let me try to rephrase, to make sure that I understood correctly.

To do this, I'd keep a single entry point mlir-to-llvmir and remove all the others (avx512-mlir-to-llvmir should not exist). The implementation LLVM::ModuleTranslation::convertOperation could be final. The implementation can simply lookup the dialect for each operation and dispatch the conversion through a new interface.

What you are suggesting is:

• There is only one entry point in the code for converting to LLVMIR: LLVM::ModuleTranslation::translateModule()
• We no longer override translateModule. No subclasses of LLVM::ModuleTranslation are needed and the entire ModuleTranslation class can be final. (However, there's one more subclass in Google3 that we have to think about.)
• ModuleTranslation::convertOperation should be able to translate all kind of operations. This can be implemented similar to how OpenMP operations are handled at the moment: We check the operation dialect and then dispatch to the corresponding convert*Operation function:

if (opInst.getDialect() == ompDialect)
  return convertOmpOperation(opInst, builder);

• All avx512-mlir-to-llvmir will be removed (for both mlir-cpu-runner and mlir-translate).

This way we don't rely on any global static registration (other than the existing MLIRContext) for the translation, new Dialect can easily plug into the LLVM translation, including out-of-tree dialects.

What is not clear to me, however, is how translateModule knows which dialects exist (as well as the pointers to their corresponding convert*Operation functions). I think we may still have to statically register each dialect that converts to LLVMIR in a map.

[mehdi_amini]

In D95732#2535596, @springerm wrote:

This way we don't rely on any global static registration (other than the existing MLIRContext) for the translation, new Dialect can easily plug into the LLVM translation, including out-of-tree dialects.

What is not clear to me, however, is how translateModule knows which dialects exist (as well as the pointers to their corresponding convert*Operation functions). I think we may still have to statically register each dialect that converts to LLVMIR in a map.

That is the part where I'd have used a Dialect Interface. Instead of:

if (opInst.getDialect() == ompDialect)
   return convertOmpOperation(opInst, builder);

We'd do:

if (auto *translateInterface = opInst.getDialect()->getRegisteredInterface<LLVMTranslateInterface>()))
   return translateInterface.convert(opInst, builder);

That way the LLVM dialect does not need to be taught about the others dialect and it can dispatch in a totally decoupled way.

The only problem here is the one @ftynse mentioned: in order to provide an implementation for the interface, the other dialects would have a build dependency to LLVM IR. Which not every MLIR clients using these dialects may want. Decoupling the interface from the dialect gets a bit more tricky.

[springerm]

In D95732#2535628, @mehdi_amini wrote:

The only problem here is the one @ftynse mentioned: in order to provide an implementation for the interface, the other dialects would have a build dependency to LLVM IR. Which not every MLIR clients using these dialects may want. Decoupling the interface from the dialect gets a bit more tricky.

I think that's already the case. There is the following dependency chain in the BUILD file: mlir-opt <- AllPassesAndDialectsNoRegistration <- LLVMDialect <- llvm:Core.

In general, it makes sense that we would want to avoid building LLVM stuff when just building mlir-opt. Of the top of my head, here's one way of separating LLVM dependencies from LLVMDialect/NVVMDialect/etc.: By having the LLVMTranslateInterface implementations in a build target separate from the dialect build target. In fact, there would be two new build targets that provide implementations of the LLVMTranslateInterface interface. One dummy target (with no LLVM dependencies) and the "real" target (with all the LLVM dependencies). mlir-opt would depend on (link) the dummy build target and mlir-translate would depend on the "real" target.

This would have to be designed carefully though, as accidentally including both dummy and real build target (or none at all) would result in a linker error.

[mehdi_amini]

In D95732#2535824, @springerm wrote:

In D95732#2535628, @mehdi_amini wrote:

The only problem here is the one @ftynse mentioned: in order to provide an implementation for the interface, the other dialects would have a build dependency to LLVM IR. Which not every MLIR clients using these dialects may want. Decoupling the interface from the dialect gets a bit more tricky.

I think that's already the case. There is the following dependency chain in the BUILD file: mlir-opt <- AllPassesAndDialectsNoRegistration <- LLVMDialect <- llvm:Core.

But only for the LLVM Dialects, I don't expect the other dialects you want to lower to LLVM to depend on LLVM Core right now.

In general, it makes sense that we would want to avoid building LLVM stuff when just building mlir-opt.

Actually mlir-opt is a testing/debugging/development tool, and the fact that it is bloated isn't a big deal. However the other libraries are intended to be able to be assembled in compilers ultimately, and preserving modularity is important :)

Of the top of my head, here's one way of separating LLVM dependencies from LLVMDialect/NVVMDialect/etc.

By having the LLVMTranslateInterface implementations in a build target separate from the dialect build target. In fact, there would be two new build targets that provide implementations of the LLVMTranslateInterface interface. One dummy target (with no LLVM dependencies) and the "real" target (with all the LLVM dependencies). mlir-opt would depend on (link) the dummy build target and mlir-translate would depend on the "real" target.

This would have to be designed carefully though, as accidentally including both dummy and real build target (or none at all) would result in a linker error.

I don't quite get how this would work though? At the moment interfaces are registered in a dialect constructor: you can build the dialect without the interfaces it defines.
We could make it so that you could register dialect interfaces separately / after the fact, but how would this work? The registration for the dialect is explicit, so would you have another function "registerAVX521TranslationInterface()" to add the interface and tools that want it would need to register the interface?

[ftynse]

But only for the LLVM Dialects, I don't expect the other dialects you want to lower to LLVM to depend on LLVM Core right now.

And only because we haven't implemented the data layout yet :) There's nothing else that we use AFAIK.

I don't quite get how this would work though? At the moment interfaces are registered in a dialect constructor: you can build the dialect without the interfaces it defines.
We could make it so that you could register dialect interfaces separately / after the fact, but how would this work? The registration for the dialect is explicit, so would you have another function "registerAVX521TranslationInterface()" to add the interface and tools that want it would need to register the interface?

I think the idea was more along the lines of the following. The interface actually dispatches to something like LogicalResult translateLLVMDialectOperation(Operation *, TranslationState &). This function is only declared in the interface. It has two definitions. One in libDummyLLVMTranslation and is "implemented" as llvm_unrecheable. The other is in libLLVMTranslation and contains the actual implementation. The build scripts link libDummyLLVMTranslation for everything but mlir-translate and similar uses.

This sounds a bit fragile to me, so I was considering some sort of dependency injection. Like the interface has a setTranslationImpl(function_ref<LogicalResult (Operation *, TranslationState &)>) (strictly speaking, this is no longer an interface because it will have to have state). By default, the translation impl is function that always fails. Clients can override it after registering dialects to the actual translation function provided by a separate library. So, in JitRunnerMain.cpp, we can have context.getOrLoadDialect<LLVMDialect>()->getRegisteredInterface<LLVMTranslationInterface>().setTranslationImpl(llvmTranslation). Here, llvmTranslation is defined in a separate library that only the clients that need it will depend on. This sounds like too much indirection though, maybe LLVM ModuleTranslation can just have a map dialect -> function and dispatch calls itself without going through interfaces.

[mehdi_amini]

In D95732#2537207, @ftynse wrote:

This sounds a bit fragile to me, so I was considering some sort of dependency injection. Like the interface has a setTranslationImpl(function_ref<LogicalResult (Operation *, TranslationState &)>) (strictly speaking, this is no longer an interface because it will have to have state). By default, the translation impl is function that always fails. Clients can override it after registering dialects to the actual translation function provided by a separate library. So, in JitRunnerMain.cpp, we can have context.getOrLoadDialect<LLVMDialect>()->getRegisteredInterface<LLVMTranslationInterface>().setTranslationImpl(llvmTranslation). Here, llvmTranslation is defined in a separate library that only the clients that need it will depend on. This sounds like too much indirection though, maybe LLVM ModuleTranslation can just have a map dialect -> function and dispatch calls itself without going through interfaces.

This requires to pre-load dialects though, which isn't desirable in general. It seems to me that these are all convoluted way to carry some "state", so why not use the existing "state" we have an register the interface on the dialect separately?

[ftynse]

In D95732#2537260, @mehdi_amini wrote:

This requires to pre-load dialects though, which isn't desirable in general. It seems to me that these are all convoluted way to carry some "state", so why not use the existing "state" we have an register the interface on the dialect separately?

Is it just a matter of making Dialect::addInterfaces public?

[mehdi_amini]

In D95732#2537288, @ftynse wrote:

In D95732#2537260, @mehdi_amini wrote:

This requires to pre-load dialects though, which isn't desirable in general. It seems to me that these are all convoluted way to carry some "state", so why not use the existing "state" we have an register the interface on the dialect separately?

Is it just a matter of making Dialect::addInterfaces public?

Yes, I'm in favor of doing this.

[mehdi_amini]

In D95732#2537293, @mehdi_amini wrote:

In D95732#2537288, @ftynse wrote:

In D95732#2537260, @mehdi_amini wrote:

This requires to pre-load dialects though, which isn't desirable in general. It seems to me that these are all convoluted way to carry some "state", so why not use the existing "state" we have an register the interface on the dialect separately?

Is it just a matter of making Dialect::addInterfaces public?

Yes, I'm in favor of doing this.

Actually scratch that, that still requires a Dialect instance, it isn't that simple.

[ftynse]

It looks like creating an interface required a dialect instance anyway, so we'd need some mechanism of delayed registration, presumably in the context, that takes a dialect TypeID and construction lambda, and registers an interface when and if the dialect is actually loaded.

[mehdi_amini]

Even this would require a context instance, which does not fit what's in mlir-cpu-runner.cpp right now:

int main(int argc, char **argv) {
  llvm::InitLLVM y(argc, argv);
  llvm::InitializeNativeTarget();
  llvm::InitializeNativeTargetAsmPrinter();
  llvm::InitializeNativeTargetAsmParser();
  mlir::initializeLLVMPasses();
  mlir::registerAllTranslations();

  return mlir::JitRunnerMain(argc, argv);
}

So we couldn't register translations in a similar fashion here I think, we'd have to go the same route as the dialect registry, with a separate "InterfaceRegistry" of some sort that can be injected into the Context the same way.

[ftynse]

Well, that does fit not because context creation is factored out into JitRunnerMain together with other code shared between between CPU and CUDA runners. I don't see a strong reason not to make JitRunnerMain take a context instead of creating it here https://github.com/llvm/llvm-project/blob/fae6d129dac27fa8c98c1680ac79af7b453ec7af/mlir/lib/ExecutionEngine/JitRunner.cpp#L333. (Also not registering all dialects, only those we actually need).

Actually, the mechanism I have in mind can be put into DialectRegistry, so we wouldn't even need a context, just the registry.

[nicolasvasilache]

I think we can drop this now?

mlir/include/mlir/Target/LLVMIR.h

 class LLVMContext;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 
 class OwningModuleRef;
 class MLIRContext;
-class ModuleOp;
+class Operation;
 
 /// Convert the given MLIR module into LLVM IR.  The LLVM context is extracted
 /// from the registered LLVM IR dialect.  In case of error, report it
 /// to the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `nullptr`.
 std::unique_ptr<llvm::Module>
-translateModuleToLLVMIR(ModuleOp m, llvm::LLVMContext &llvmContext,
+translateModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
                         StringRef name = "LLVMDialectModule");
 
 /// Convert the given LLVM module into MLIR's LLVM dialect.  The LLVM context is
 /// extracted from the registered LLVM IR dialect. In case of error, report it
 /// to the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `{}`.
 OwningModuleRef
 translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                         MLIRContext *context);
 
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_H

mlir/include/mlir/Translation.h

 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_TRANSLATION_H
 #define MLIR_TRANSLATION_H
 
 #include "llvm/Support/CommandLine.h"
 
 namespace llvm {
+class LLVMContext;
 class MemoryBuffer;
+class Module;
 class SourceMgr;
 class StringRef;
 } // namespace llvm
 
 namespace mlir {
 class DialectRegistry;
 struct LogicalResult;
 class MLIRContext;
 class ModuleOp;
+class Operation;
 class OwningModuleRef;
 
 /// Interface of the function that translates the sources managed by `sourceMgr`
 /// to MLIR. The source manager has at least one buffer. The implementation
 /// should create a new MLIR ModuleOp in the given context and return a pointer
 /// to it, or a nullptr in case of any error.
 using TranslateSourceMgrToMLIRFunction =
     std::function<OwningModuleRef(llvm::SourceMgr &sourceMgr, MLIRContext *)>;

 /// 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. 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(
     llvm::SourceMgr &sourceMgr, llvm::raw_ostream &output, MLIRContext *)>;
 
+/// Interface of the function that translates an MLIR module into an LLVM
+/// module. Such functions can be passed to mlir-cpu-runner to generate LLVM IR
+/// for a specific target architecture.
+using TranslateMLIRToLLVMFunction = std::function<std::unique_ptr<llvm::Module>(
+    Operation *, llvm::LLVMContext &)>;
+
 /// Use Translate[ToMLIR|FromMLIR]Registration as an initializer that
 /// registers a function and associates it with name. This requires that a
 /// translation has not been registered to a given name.
 ///
+/// Use TranslateFromMLIRToLLVMRegistration for MLIR to LLVMIR translations.
+/// Such translations are made available to mlir-cpu-runner and can be
+/// specified via command line parameter.
+///
 /// Usage:
 ///
 ///   // At file scope.
 ///   namespace mlir {
 ///   void registerTRexToMLIRRegistration() {
 ///     TranslateToMLIRRegistration Unused(&MySubCommand, [] { ... });
 ///   }
 ///   } // namespace mlir
 ///
 /// \{
 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,
       std::function<void(DialectRegistry &)> dialectRegistration =
           [](DialectRegistry &) {});
 };
+struct TranslateFromMLIRToLLVMRegistration
+    : public TranslateFromMLIRRegistration {
+  TranslateFromMLIRToLLVMRegistration(
+      llvm::StringRef name, const TranslateMLIRToLLVMFunction &function,
+      std::function<void(DialectRegistry &)> dialectRegistration =
+          [](DialectRegistry &) {});
+};
 struct TranslateRegistration {
   TranslateRegistration(llvm::StringRef name,
                         const TranslateFunction &function);
 };
 /// \}
 
-/// A command line parser for translation functions.
-struct TranslationParser : public llvm::cl::parser<const TranslateFunction *> {
-  TranslationParser(llvm::cl::Option &opt);
+/// A command line parser for translation functions. F specifies the type of
+/// translations: TranslateFunction or TranslateMLIRToLLVMFunction.
+template <typename F>
+struct TranslationParserBase : public llvm::cl::parser<const F *> {
+  TranslationParserBase(llvm::cl::Option &opt);
 
   void printOptionInfo(const llvm::cl::Option &o,
                        size_t globalWidth) const override;
 };
 
+/// A command line parser for general translation functions.
+using TranslationParser = TranslationParserBase<TranslateFunction>;
+
+/// A command line parser for MLIR to LLVMIR translations to be used with
+/// mlir-cpu-runner.
+using MLIRToLLVMTranslationParser =
+    TranslationParserBase<TranslateMLIRToLLVMFunction>;
+
 /// Translate to/from an MLIR module from/to an external representation (e.g.
 /// LLVM IR, SPIRV binary, ...). This is the entry point for the implementation
 /// of tools like `mlir-translate`. The translation to perform is parsed from
 /// the command line. The `toolName` argument is used for the header displayed
 /// by `--help`.
 LogicalResult mlirTranslateMain(int argc, char **argv,
                                 llvm::StringRef toolName);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSLATION_H

mlir/lib/ExecutionEngine/JitRunner.cpp

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
+#include "mlir/Translation.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
struct Options {
   llvm::cl::opt<bool> dumpObjectFile{
       "dump-object-file",
       llvm::cl::desc("Dump JITted-compiled object to file specified with "
                      "-object-filename (<input file>.o by default).")};
 
   llvm::cl::opt<std::string> objectFilename{
       "object-filename",
       llvm::cl::desc("Dump JITted-compiled object to file <input file>.o")};
+
+  // Add flags for all the registered MLIR to LLVMIR translations.
+  llvm::cl::opt<const TranslateMLIRToLLVMFunction *, false,
+                MLIRToLLVMTranslationParser>
+      translationRequested{"translation",
+                           llvm::cl::desc("Translation to perform")};
 };
 
 struct CompileAndExecuteConfig {
   /// LLVM module transformer that is passed to ExecutionEngine.
   llvm::function_ref<llvm::Error(llvm::Module *)> transformer;
 
   /// A custom function that is passed to ExecutionEngine. It processes MLIR
   /// module and creates LLVM IR module.
int mlir::JitRunnerMain(int argc, char **argv, JitRunnerConfig config) {
   auto transformer = mlir::makeLLVMPassesTransformer(
       passes, optLevel, /*targetMachine=*/tmOrError->get(), optPosition);
 
   CompileAndExecuteConfig compileAndExecuteConfig;
   compileAndExecuteConfig.transformer = transformer;
   compileAndExecuteConfig.llvmModuleBuilder = config.llvmModuleBuilder;
   compileAndExecuteConfig.runtimeSymbolMap = config.runtimesymbolMap;
 
+  if (const auto *f = options.translationRequested.getValue()) {
+    assert(!compileAndExecuteConfig.llvmModuleBuilder &&
+           "Competing translations specified via both command line and config");
+    compileAndExecuteConfig.llvmModuleBuilder = *f;
+  }
+
   // Get the function used to compile and execute the module.
   using CompileAndExecuteFnT =
       Error (*)(Options &, ModuleOp, StringRef, CompileAndExecuteConfig);
   auto compileAndExecuteFn =
       StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
           .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
           .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
           .Case("f32", compileAndExecuteSingleReturnFunction<float>)

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 #include "mlir/Translation.h"
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/ToolOutputFile.h"
 
-using namespace mlir;
-
+namespace mlir {
 std::unique_ptr<llvm::Module>
-mlir::translateModuleToLLVMIR(ModuleOp m, llvm::LLVMContext &llvmContext,
-                              StringRef name) {
+translateModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
   auto llvmModule =
       LLVM::ModuleTranslation::translateModule<>(m, llvmContext, name);
   if (!llvmModule)
-    emitError(m.getLoc(), "Fail to convert MLIR to LLVM IR");
+    emitError(m->getLoc(), "Fail to convert MLIR to LLVM IR");
   else if (verifyModule(*llvmModule))
-    emitError(m.getLoc(), "LLVM IR fails to verify");
+    emitError(m->getLoc(), "LLVM IR fails to verify");
   return llvmModule;
 }
 
-namespace mlir {
 void registerToLLVMIRTranslation() {
-  TranslateFromMLIRRegistration registration(
+  TranslateFromMLIRToLLVMRegistration(
       "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();
+      [](Operation *m, llvm::LLVMContext &llvmContext) {
+        return translateModuleToLLVMIR(m, llvmContext, "LLVMDialectModule");
       },
       [](DialectRegistry &registry) {
         registry.insert<LLVM::LLVMDialect, omp::OpenMPDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp

llvmModule->getOrInsertNamedMetadata("nvvm.annotations")
         ->addOperand(llvmMetadataNode);
   }
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToNVVMIRTranslation() {
-  TranslateFromMLIRRegistration registration(
+  TranslateFromMLIRToLLVMRegistration registration(
       "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();
+      [](Operation *m, llvm::LLVMContext &llvmContext) {
+        return translateModuleToNVVMIR(m, llvmContext);
       },
       [](DialectRegistry &registry) {
         registry.insert<LLVM::LLVMDialect, NVVM::NVVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

for (auto func :
     llvmFunc->addFnAttr("amdgpu-flat-work-group-size", "1, 1024");
   }
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToROCDLIRTranslation() {
-  TranslateFromMLIRRegistration registration(
+  TranslateFromMLIRToLLVMRegistration registration(
       "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();
+      [](Operation *m, llvm::LLVMContext &llvmContext) {
+        return translateModuleToROCDLIR(m, llvmContext);
       },
       [](DialectRegistry &registry) {
         registry.insert<ROCDL::ROCDLDialect, LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp

 protected:
   LogicalResult convertOperation(Operation &opInst,
                                  llvm::IRBuilder<> &builder) override {
 #include "mlir/Dialect/LLVMIR/LLVMAVX512Conversions.inc"
 
     return LLVM::ModuleTranslation::convertOperation(opInst, builder);
   }
 };
-
-std::unique_ptr<llvm::Module>
-translateLLVMAVX512ModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
-                                  StringRef name) {
-  return LLVM::ModuleTranslation::translateModule<LLVMAVX512ModuleTranslation>(
-      m, llvmContext, name);
-}
 } // end namespace
 
 namespace mlir {
 void registerAVX512ToLLVMIRTranslation() {
-  TranslateFromMLIRRegistration reg(
+  TranslateFromMLIRToLLVMRegistration(
       "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();
+      [](Operation *m, llvm::LLVMContext &llvmContext) {
+        return LLVM::ModuleTranslation::translateModule<
+            LLVMAVX512ModuleTranslation>(m, llvmContext, "LLVMDialectModule");
       },
       [](DialectRegistry &registry) {
         registry.insert<LLVM::LLVMAVX512Dialect, LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/LLVMArmNeonIntr.cpp

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

mlir/lib/Target/LLVMIR/LLVMArmSVEIntr.cpp

 translateLLVMArmSVEModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
                                   StringRef name) {
   return LLVM::ModuleTranslation::translateModule<LLVMArmSVEModuleTranslation>(
       m, llvmContext, name);
 }
 
 namespace mlir {
 void registerArmSVEToLLVMIRTranslation() {
-  TranslateFromMLIRRegistration reg(
+  TranslateFromMLIRToLLVMRegistration reg(
       "arm-sve-mlir-to-llvmir",
-      [](ModuleOp module, raw_ostream &output) {
-        llvm::LLVMContext llvmContext;
-        auto llvmModule = translateLLVMArmSVEModuleToLLVMIR(
-            module, llvmContext, "LLVMDialectModule");
-        if (!llvmModule)
-          return failure();
-
-        llvmModule->print(output, nullptr);
-        return success();
+      [](Operation *m, llvm::LLVMContext &llvmContext) {
+        return translateLLVMArmSVEModuleToLLVMIR(m, llvmContext,
+                                                 "LLVMDialectModule");
       },
       [](DialectRegistry &registry) {
         registry.insert<LLVM::LLVMArmSVEDialect, LLVM::LLVMDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Translation/Translation.cpp

 
 #include "mlir/Translation.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/ToolUtilities.h"
+#include "mlir/Target/LLVMIR/ModuleTranslation.h"
+
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // 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;
+/// A TranslationRegistry stores mappings betwen MLIR translations and the
+/// functions that perform those translations. There are two kind of
+/// translations:
+///
+/// * File-to-file MLIR translations (F = TranslateFunctions)
+/// * MLIR module to LLVM module translations (F = TranslateMLIRToLLVMFunction)
+///
+/// Only the second kind of translations can be used with mlir-cpu-runner.
+template <typename F>
+struct TranslationRegistry {
+  /// Get the mutable static map between registered MLIR translations and the
+  /// functions that perform those translations.
+  static llvm::StringMap<F> &get() {
+    static llvm::StringMap<F> translationRegistry;
     return translationRegistry;
   }
 
   /// Register the given translation.
-static void registerTranslation(StringRef name,
-                                const TranslateFunction &function) {
-  auto &translationRegistry = getTranslationRegistry();
+  static void registerTranslation(StringRef name, const F &function) {
+    auto &translationRegistry = get();
     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);
+  TranslationRegistry<TranslateFunction>::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 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);
+  TranslationRegistry<TranslateFunction>::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

 
 //===----------------------------------------------------------------------===//
 // Translation from MLIR
 //===----------------------------------------------------------------------===//
 
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, const TranslateFromMLIRFunction &function,
     std::function<void(DialectRegistry &)> dialectRegistration) {
-  registerTranslation(name, [function, dialectRegistration](
-                                llvm::SourceMgr &sourceMgr, raw_ostream &output,
+  auto fn = [function, dialectRegistration](llvm::SourceMgr &sourceMgr,
+                                            raw_ostream &output,
                                             MLIRContext *context) {
     dialectRegistration(context->getDialectRegistry());
     auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
     if (!module)
       return failure();
     return function(module.get(), output);
-  });
+  };
+  TranslationRegistry<TranslateFunction>::registerTranslation(name, fn);
+}
+
+//===----------------------------------------------------------------------===//
+// Translation from MLIR to LLVM IR
+//===----------------------------------------------------------------------===//
+
+TranslateFromMLIRToLLVMRegistration::TranslateFromMLIRToLLVMRegistration(
+    StringRef name, const TranslateMLIRToLLVMFunction &function,
+    std::function<void(DialectRegistry &)> dialectRegistration)
+    : TranslateFromMLIRRegistration(
+          /*name=*/name,
+          /*function=*/
+          [function, dialectRegistration](ModuleOp module,
+                                          llvm::raw_ostream &output) {
+            llvm::LLVMContext llvmContext;
+            auto llvmModule = function(module, llvmContext);
+            if (!llvmModule)
+              return failure();
+
+            llvmModule->print(output, nullptr);
+            return success();
+          },
+          /*dialectRegistration=*/dialectRegistration) {
+  TranslationRegistry<TranslateMLIRToLLVMFunction>::registerTranslation(
+      name, function);
 }
 
 //===----------------------------------------------------------------------===//
 // Translation Parser
 //===----------------------------------------------------------------------===//
 
-TranslationParser::TranslationParser(llvm::cl::Option &opt)
-    : llvm::cl::parser<const TranslateFunction *>(opt) {
-  for (const auto &kv : getTranslationRegistry())
-    addLiteralOption(kv.first(), &kv.second, kv.first());
+template <typename F>
+TranslationParserBase<F>::TranslationParserBase(llvm::cl::Option &opt)
+    : llvm::cl::parser<const F *>(opt) {
+  for (const auto &kv : TranslationRegistry<F>::get())
+    this->addLiteralOption(kv.first(), &kv.second, kv.first());
 }
 
-void TranslationParser::printOptionInfo(const llvm::cl::Option &o,
+template <typename F>
+void TranslationParserBase<F>::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);
+  TranslationParserBase<F> *tp = const_cast<TranslationParserBase<F> *>(this);
+  llvm::array_pod_sort(
+      tp->Values.begin(), tp->Values.end(),
+      [](const typename TranslationParserBase<F>::OptionInfo *lhs,
+         const typename TranslationParserBase<F>::OptionInfo *rhs) {
+  llvm::cl::parser<const F *>::printOptionInfo(o, globalWidth);
 }
 
+//===----------------------------------------------------------------------===//
+// Explicit Translation Parser Instantiations
+//===----------------------------------------------------------------------===//
+
+template class mlir::TranslationParserBase<TranslateFunction>;
+template class mlir::TranslationParserBase<TranslateMLIRToLLVMFunction>;
+
 LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
                                       llvm::StringRef toolName) {
 
   static llvm::cl::opt<std::string> inputFilename(
       llvm::cl::Positional, llvm::cl::desc("<input file>"),
       llvm::cl::init("-"));
 
   static llvm::cl::opt<std::string> outputFilename(

mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp

 // CPU by  translating MLIR to LLVM IR before JIT-compiling and executing the
 // latter.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllTranslations.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   llvm::InitializeNativeTargetAsmParser();
   mlir::initializeLLVMPasses();
+  mlir::registerAllTranslations();
 
   return mlir::JitRunnerMain(argc, argv);
 }