[MLIR] Add convertModule functionallity to the LLVMTranslationInterface and ModuleTranslation

[agozillon]

Summary

This patch seeks to add a convertModuleOperation function which allows
dialect specific module operation lowering to LLVM-IR (the operation
itself, not its body and for a dialect module, e.g. an OpenMP module
such as omp.module). It also allows amendOperation to be invoked on
module attributes for builtin and specialised modules, allowing
dialect specific attribute lowering, e.g. I add an omp.is_device flag,
which can now have specialised lowering to LLVM IR.

To give some context, I have currently found use for this in the
OpenMPIRToLLVMTranslation lowering segment (utilised in the
Flang compiler). It allows the ability to hook into module attributes
that are specific to the OpenMP Dialect that need lowering using
the OpenMPIRBuilder that is shared between Clang/Flang. I have
previously tested it with an OpenMP specific module.
Although, I am aiming to upstream the former, rather than the
latter if this is the correct approach. I have an external patch
that I can link to which I am trying to upstream in components,
if that helps to give more context.

[ftynse]

Why does it have to be a separate interface function? Module is just the builtin.module operation, it looks like we can just call convertOperation for it and teach the default translation to do nothing with it.

[agozillon]

In D144883#4156106, @ftynse wrote:

Why does it have to be a separate interface function? Module is just the builtin.module operation, it looks like we can just call convertOperation for it and teach the default translation to do nothing with it.

I could definitely try this approach, I haven't yet as I thought it would be best to keep them distinct for explicitness's sake and in-case there was some digression between them but considering they are identical at the moment I don't see why they couldn't be merged, thank you for the suggestion I'll see if it is possible to clean the patch up that way!

[agozillon]

In D144883#4156106, @ftynse wrote:

Why does it have to be a separate interface function? Module is just the builtin.module operation, it looks like we can just call convertOperation for it and teach the default translation to do nothing with it.

I found out the reason (again possibly) I didn't do it this way! The convertModuleOperation works slightly different to convertOperation, convertModuleOperation will return success rather than failure when it encounters an operation (module) that has no interface for lowering it, allowing other dialects the option of lowering the module or not lowering the module as the 'builtin.module' does. This allows existing dialect modules such as gpu.module to continue to not do dialect specific lowering, allowing them to opt in or out based on their dialects needs.

[ftynse]

This allows existing dialect modules such as gpu.module to continue to not do dialect specific lowering, allowing them to opt in or out based on their dialects needs.

This is why I also mentioned that we need to "teach the default translation to do nothing with it". Can't we add the interface implementation that returns success instead of hardcoding it as default?

[agozillon]

In D144883#4158684, @ftynse wrote:

This allows existing dialect modules such as gpu.module to continue to not do dialect specific lowering, allowing them to opt in or out based on their dialects needs.

This is why I also mentioned that we need to "teach the default translation to do nothing with it". Can't we add the interface implementation that returns success instead of hardcoding it as default?

In this case which interface implementation do you mean? Sorry for the dumb question, I will likely have one or two others tomorrow after I've had a chance to have a deeper look into this change request!

What I have currently tried to do was remove the convertModuleOperation related changes, invoke convertOperation on the Module in the same location and wrap part of the body of the convertOperation function belonging to ModuleTranslation.cpp in a check which tests if the operation is a builtin.module, if it is it would skip the searching for a dialect interface and the forwarding on to the convertOperation for that dialect interface (which would prevent error emission from not finding a dialect interface for the builtin.module). This however, still caused tests to fail for the gpu.module as the module wasn't handled by it's interface as of yet. This change very likely isn't what you mean, so I am perhaps missing a piece of knowledge to get to the solution you're suggesting here, so I apologies if I ask some silly questions while trying to get further clarity on it!

[agozillon]

After a little bit of time going over the suggested change of ditching convertModuleOperation and using the existing convertOperation infrastructure, I can see two ways of doing it:

• Teach the default LLVMTranslationInterface::convertOperation to ignore builtin.module and gpu.module and return success, but this also means removing the error output for not finding a dialect interface for an operation inside of the ModuleTranslation::convertOperation function, otherwise we can hoist this ignore from LLVMTranslationInterface::convertOperation to ModuleTranslation::convertOperation and keep the error emission.
• Create and register a LLVMTranslationInterface for gpu and builtin dialects that return success and do nothing else.

I think you were meaning the latter, I am happy to do either, or even another alternative if I have misunderstood what you mean!

[ftynse]

In D144883#4161426, @agozillon wrote:

After a little bit of time going over the suggested change of ditching convertModuleOperation and using the existing convertOperation infrastructure, I can see two ways of doing it:

• Teach the default LLVMTranslationInterface::convertOperation to ignore builtin.module and gpu.module and return success, but this also means removing the error output for not finding a dialect interface for an operation inside of the ModuleTranslation::convertOperation function, otherwise we can hoist this ignore from LLVMTranslationInterface::convertOperation to ModuleTranslation::convertOperation and keep the error emission.
• Create and register a LLVMTranslationInterface for gpu and builtin dialects that return success and do nothing else.

I think you were meaning the latter, I am happy to do either, or even another alternative if I have misunderstood what you mean!

I indeed meant the latter. We can also use the external model mechanism https://mlir.llvm.org/docs/Interfaces/#external-models-for-attribute-operation-and-type-interfaces so we don't need to make the builtin dialect aware of the LLVM translation interface.

[agozillon]

Going to abandon and close this revision as the Reviewers point was implemented and accepted in the following patch https://reviews.llvm.org/D145932/ by @skatrak, thank you very much @skatrak for looking into this! And thank you @ftynse for recommending a better direction.

mlir/include/mlir/Target/LLVMIR/LLVMTranslationInterface.h

public:
   /// Hook for derived dialect interface to provide translation of the
   /// operations to LLVM IR.
   virtual LogicalResult
   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) const {
     return failure();
   }
 
+  /// Hook for derived dialect interface to provide translation of module
+  /// type operations to LLVM IR where special handling is neccessary.
+  virtual LogicalResult
+  convertModuleOperation(Operation *op, llvm::IRBuilderBase &builder,
+                         LLVM::ModuleTranslation &moduleTranslation) const {
+    return success();
+  }
+
   /// Hook for derived dialect interface to act on an operation that has dialect
   /// attributes from the derived dialect (the operation itself may be from a
   /// different dialect). This gets called after the operation has been
   /// translated. The hook is expected to use moduleTranslation to look up the
   /// translation results and amend the corresponding IR constructs. Does
   /// nothing and succeeds by default.
   virtual LogicalResult
   amendOperation(Operation *op, NamedAttribute attribute,
public:
   virtual LogicalResult
   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) const {
     if (const LLVMTranslationDialectInterface *iface = getInterfaceFor(op))
       return iface->convertOperation(op, builder, moduleTranslation);
     return failure();
   }
 
+  /// Translates the given module operation to LLVM IR using the interface
+  /// implemented by the op's dialect. This is invoked after all operations
+  /// within the module have been processed.
+  virtual LogicalResult
+  convertModuleOperation(Operation *op, llvm::IRBuilderBase &builder,
+                         LLVM::ModuleTranslation &moduleTranslation) const {
+    if (const LLVMTranslationDialectInterface *iface = getInterfaceFor(op))
+      return iface->convertModuleOperation(op, builder, moduleTranslation);
+    return success();
+  }
+
   /// Acts on the given operation using the interface implemented by the dialect
   /// of one of the operation's dialect attributes.
   virtual LogicalResult
   amendOperation(Operation *op, NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const {
     if (const LLVMTranslationDialectInterface *iface =
             getInterfaceFor(attribute.getNameDialect())) {
       return iface->amendOperation(op, attribute, moduleTranslation);
     }
     return success();
   }
 };
 
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_LLVMTRANSLATIONINTERFACE_H

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

private:
 
   /// Converts individual components.
   LogicalResult convertOperation(Operation &op, llvm::IRBuilderBase &builder);
   LogicalResult convertFunctionSignatures();
   LogicalResult convertFunctions();
   LogicalResult convertGlobals();
   LogicalResult convertOneFunction(LLVMFuncOp func);
 
+  /// Converts a module operation after all other operations have been handled.
+  /// This is a No-Op for dialects that have no specified interface for module
+  /// operations or attributes. This function is invoked after all operations
+  /// in the module have been processed.
+  LogicalResult convertModuleOperation(Operation &op,
+                                       llvm::IRBuilderBase &builder);
+
   /// Process access_group LLVM Metadata operations and create LLVM
   /// metadata nodes.
   LogicalResult createAccessGroupMetadata();
 
   /// Process alias.scope LLVM Metadata operations and create LLVM
   /// metadata nodes for them and their domains.
   LogicalResult createAliasScopeMetadata();
 

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 llvm::Value *mlir::LLVM::detail::createIntrinsicCall(
     llvm::IRBuilderBase &builder, llvm::Intrinsic::ID intrinsic,
     ArrayRef<llvm::Value *> args, ArrayRef<llvm::Type *> tys) {
   llvm::Module *module = builder.GetInsertBlock()->getModule();
   llvm::Function *fn = llvm::Intrinsic::getDeclaration(module, intrinsic, tys);
   return builder.CreateCall(fn, args);
 }
 
+/// Given a single MLIR module, create the corresponding LLVM IR operations
+/// using the `builder`, this lowering is defined by the dialect and for dialect
+/// specific modules e.g. omp.module from the OpenMP dialect. This is a
+/// completely optional hook for a dialect. This function is invoked after all
+/// operations in the module have been processed.
+LogicalResult
+ModuleTranslation::convertModuleOperation(Operation &op,
+                                          llvm::IRBuilderBase &builder) {
+  if (const LLVMTranslationDialectInterface *opIface =
+          iface.getInterfaceFor(&op)) {
+    if (failed(opIface->convertModuleOperation(&op, builder, *this)))
+      return failure();
+  }
+  return convertDialectAttributes(&op);
+}
+
 /// Given a single MLIR operation, create the corresponding LLVM IR operation
 /// using the `builder`.
 LogicalResult
 ModuleTranslation::convertOperation(Operation &op,
                                     llvm::IRBuilderBase &builder) {
   const LLVMTranslationDialectInterface *opIface = iface.getInterfaceFor(&op);
   if (!opIface)
     return op.emitError("cannot be converted to LLVM IR: missing "
for (Operation &o : getModuleBody(module).getOperations()) {
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
              LLVM::GlobalDtorsOp, LLVM::MetadataOp>(&o) &&
         !o.hasTrait<OpTrait::IsTerminator>() &&
         failed(translator.convertOperation(o, llvmBuilder))) {
       return nullptr;
     }
   }
 
+  // Convert a dialect specific module and it's attributes with dialect
+  // specific lowering if neccessary, otherwise this is a no-op
+  if (failed(translator.convertModuleOperation(*module, llvmBuilder)))
+    return nullptr;
+
   if (llvm::verifyModule(*translator.llvmModule, &llvm::errs()))
     return nullptr;
 
   return std::move(translator.llvmModule);
 }