[MLIR] Move builtin.module LLVM IR translation to before nested operations

[skatrak]

Summary

This patch moves the call for translating an MLIR module to LLVM IR to the beginning of the translation process. This enables the use of dialect attributes attached to builtin.module operations and the amendOperation() flow to initialize dialect-specific global configuration before translating the contents of the module.

Currently, this patch does not impact the generated IR on its own.

[skatrak]

Ping for review! Thank you

[kiranchandramohan]

LGTM. Please wait a couple of days to see if @ftynse @gysit @Dinistro @zero9178 have any comments here.

[zero9178]

Makes sense to me. IIRC the dialect conversion framework also does things in pre-order. This now nicely turns the translation into "module -> top-level-ops -> function-bodies" order.

[gysit]

The change looks good to me!

Would it be possible to add a test, e.g. by implementing the translation interface for the test dialect and then attaching a test dialect attribute to the module?

There have been other revisions that modified the translation order and it may be a good idea to have some testing to ensure follow up commits do not break previous assumptions.

[skatrak]

Add testing infrastructure to allow translating operations from the Test dialect to LLVM IR, and add a simple test to make sure that module attributes are translated before the module body.

[skatrak]

Thank you all for having a look and for your suggestions! They should be addressed now, but I'd like to request another quick review, since I had to make significant additions in order to test this change using the Test dialect.

[zero9178]

mlir/test/lib/Dialect/Test/CMakeLists.txt

87–102

Having the TestDialect depend on these is problematic I believe as these depend on LLVM proper so will pull in LLVM IR and such as link constructs.
You'd likely see this if building with shared libs.

I think the best way to implement this is to make a separate add_mlir_translation_library just for TestToLLVMIRTranslation that depends on TestDialect, builds`TestToLLVMIRTranslation.cpp` and make mlir-translate link against that if MLIR_INCLUDE_TESTS is true in CMake.

[gysit]

Thanks for adding the test!

Looks great modulo the packaging in a separate library suggested by @zero9178.

[skatrak]

Split translation of Test dialect to LLVM IR into its own MLIR translation library.

[skatrak]

mlir/test/lib/Dialect/Test/CMakeLists.txt

87–102

Thank you for the pointers, I just finished making these changes. In the end, I was able to keep the CMakeLists.txt for mlir-translate unmodified because the add_mlir_translation_library call registers the library into the MLIR_TRANSLATION_LIBS variable that's already used by mlir-translate to find libraries to link to.

[zero9178]

LGTM

[gysit]

Thank you for the pointers, I just finished making these changes. In the end, I was able to keep the CMakeLists.txt for mlir-translate unmodified because the add_mlir_translation_library call registers the library into the MLIR_TRANSLATION_LIBS variable that's already used by mlir-translate to find libraries to link to.

Nice!

LGTM

[skatrak]

This seems to have broken post-commit builds (1, 2), so I reverted the change until I can take a look at it next week. Sorry for the noise!

[zero9178]

In D158278#4634775, @skatrak wrote:

This seems to have broken post-commit builds (1, 2), so I reverted the change until I can take a look at it next week. Sorry for the noise!

No worries! Reverting is the right thing to do :) I believe you might have to use EXCLUDE_FROM_LIBMLIR in your add_mlir_translation_library, just like the test dialect does above. The windows build failure seems unrelated to your change.

[skatrak]

In D158278#4634799, @zero9178 wrote:

In D158278#4634775, @skatrak wrote:

This seems to have broken post-commit builds (1, 2), so I reverted the change until I can take a look at it next week. Sorry for the noise!

No worries! Reverting is the right thing to do :) I believe you might have to use EXCLUDE_FROM_LIBMLIR in your add_mlir_translation_library, just like the test dialect does above. The windows build failure seems unrelated to your change.

Thank you, that does seem to be the solution to the problem. I should have read your comment before spending the time reaching the same conclusion! I'll re-land this shortly.

[mehdi_amini]

Please always document the commit message when doing reverts including why a revert is happening.

Also instead when reverting reverts, include the original commit message (and mention "re-land commit abc12345" in the description instead), otherwise git log / git blame is harder for everyone.

[skatrak]

In D158278#4642166, @mehdi_amini wrote:

Please always document the commit message when doing reverts including why a revert is happening.

Also instead when reverting reverts, include the original commit message (and mention "re-land commit abc12345" in the description instead), otherwise git log / git blame is harder for everyone.

Yes, sorry about that. I did get similar feedback right after I re-landed the patch, I'll keep this in mind for next time.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
   std::unique_ptr<llvm::Module> llvmModule =
       prepareLLVMModule(module, llvmContext, name);
   if (!llvmModule)
     return nullptr;
 
   LLVM::ensureDistinctSuccessors(module);
 
   ModuleTranslation translator(module, std::move(llvmModule));
+  llvm::IRBuilder<> llvmBuilder(llvmContext);
+
+  // Convert module before functions and operations inside, so dialect
+  // attributes can be used to change dialect-specific global configurations via
+  // `amendOperation()`. These configurations can then influence the translation
+  // of operations afterwards.
+  if (failed(translator.convertOperation(*module, llvmBuilder)))
+    return nullptr;
+
   if (failed(translator.convertComdats()))
     return nullptr;
   if (failed(translator.convertFunctionSignatures()))
     return nullptr;
   if (failed(translator.convertGlobals()))
     return nullptr;
   if (failed(translator.createTBAAMetadata()))
     return nullptr;
 
   // Convert other top-level operations if possible.
-  llvm::IRBuilder<> llvmBuilder(llvmContext);
   for (Operation &o : getModuleBody(module).getOperations()) {
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
              LLVM::GlobalDtorsOp, LLVM::ComdatOp>(&o) &&
         !o.hasTrait<OpTrait::IsTerminator>() &&
         failed(translator.convertOperation(o, llvmBuilder))) {
       return nullptr;
     }
   }
 
   // Operations in function bodies with symbolic references must be converted
   // after the top-level operations they refer to are declared, so we do it
   // last.
   if (failed(translator.convertFunctions()))
     return nullptr;
 
-  // Convert module itself.
-  if (failed(translator.convertOperation(*module, llvmBuilder)))
-    return nullptr;
-
   if (llvm::verifyModule(*translator.llvmModule, &llvm::errs()))
     return nullptr;
 
   return std::move(translator.llvmModule);
 }

mlir/test/Target/LLVMIR/test.mlir

+// RUN: mlir-translate -test-to-llvmir -split-input-file %s | FileCheck %s
+
+module {
+  "test.symbol"() <{sym_name = "foo"}> : () -> ()
+}
+
+// CHECK-NOT: @sym_from_attr
+// CHECK: @foo = external global i32
+// CHECK-NOT: @sym_from_attr
+
+// -----
+
+// Make sure that the module attribute is processed before its body, so that the
+// `test.symbol` that is created as a result of the `test.discardable_mod_attr`
+// attribute is later picked up and translated to LLVM IR.
+module attributes {test.discardable_mod_attr = true} {}
+
+// CHECK: @sym_from_attr = external global i32

mlir/test/lib/Dialect/Test/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   TestDialect.cpp
   TestPatterns.cpp
   TestTraits.cpp
+  TestToLLVMIRTranslation.cpp
 )
 
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
 mlir_tablegen(TestAttrInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(TestAttrInterfaces.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(TestTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TestTypeInterfaces.cpp.inc -gen-type-interface-defs)
 mlir_tablegen(TestOpInterfaces.h.inc -gen-op-interface-decls)
add_mlir_library(MLIRTestDialect
   MLIRLinalgDialect
   MLIRLinalgTransforms
   MLIRLLVMDialect
   MLIRPass
   MLIRReduce
   MLIRTensorDialect
   MLIRTransformUtils
   MLIRTransforms
   )
+
+add_mlir_translation_library(MLIRTestToLLVMIRTranslation
+  TestToLLVMIRTranslation.cpp
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLLVMDialect
+  MLIRTestDialect
+  MLIRSupport
+  MLIRBuiltinToLLVMIRTranslation
+  MLIRLLVMToLLVMIRTranslation
+  )

mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp

+//===- TestToLLVMIRTranslation.cpp - Translate Test dialect to LLVM IR ----===//
+//
+// 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 between the MLIR Test dialect and LLVM IR.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/LLVMTranslationInterface.h"
+#include "mlir/Target/LLVMIR/ModuleTranslation.h"
+#include "mlir/Tools/mlir-translate/Translation.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+namespace {
+
+class TestDialectLLVMIRTranslationInterface
+    : public LLVMTranslationDialectInterface {
+public:
+  using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
+
+  LogicalResult
+  amendOperation(Operation *op, NamedAttribute attribute,
+                 LLVM::ModuleTranslation &moduleTranslation) const final;
+
+  LogicalResult
+  convertOperation(Operation *op, llvm::IRBuilderBase &builder,
+                   LLVM::ModuleTranslation &moduleTranslation) const final;
+};
+
+} // namespace
+
+LogicalResult TestDialectLLVMIRTranslationInterface::amendOperation(
+    Operation *op, NamedAttribute attribute,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+  return llvm::StringSwitch<llvm::function_ref<LogicalResult(Attribute)>>(
+             attribute.getName())
+      // The `test.discardable_mod_attr` attribute, if present and set to
+      // `true`, results in the addition of a `test.symbol` in the module it is
+      // attached to with name "sym_from_attr".
+      .Case("test.discardable_mod_attr",
+            [&](Attribute attr) {
+              if (!isa<ModuleOp>(op)) {
+                op->emitOpError("attribute 'test.discardable_mod_attr' only "
+                                "supported in modules");
+                return failure();
+              }
+
+              bool createSymbol = false;
+              if (auto boolAttr = attr.dyn_cast<BoolAttr>())
+                createSymbol = boolAttr.getValue();
+
+              if (createSymbol) {
+                OpBuilder builder(op->getRegion(0));
+                builder.create<test::SymbolOp>(
+                    op->getLoc(),
+                    StringAttr::get(op->getContext(), "sym_from_attr"),
+                    /*sym_visibility=*/nullptr);
+              }
+
+              return success();
+            })
+      .Default([](Attribute) {
+        // Skip other discardable dialect attributes.
+        return success();
+      })(attribute.getValue());
+}
+
+LogicalResult TestDialectLLVMIRTranslationInterface::convertOperation(
+    Operation *op, llvm::IRBuilderBase &builder,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+  return llvm::TypeSwitch<Operation *, LogicalResult>(op)
+      // `test.symbol`s are translated into global integers in LLVM IR, with a
+      // name equal to the symbol they are translated from.
+      .Case([&](test::SymbolOp symOp) {
+        llvm::Module *mod = moduleTranslation.getLLVMModule();
+        llvm::IntegerType *i32Type =
+            llvm::IntegerType::get(moduleTranslation.getLLVMContext(), 32);
+        mod->getOrInsertGlobal(symOp.getSymName(), i32Type);
+        return success();
+      })
+      .Default([&](Operation *) {
+        return op->emitOpError("unsupported translation of test operation");
+      });
+}
+
+namespace mlir {
+
+void registerTestToLLVMIR() {
+  TranslateFromMLIRRegistration registration(
+      "test-to-llvmir", "test dialect to LLVM IR",
+      [](Operation *op, raw_ostream &output) {
+        llvm::LLVMContext llvmContext;
+        auto llvmModule = translateModuleToLLVMIR(op, llvmContext);
+        if (!llvmModule)
+          return failure();
+
+        llvmModule->print(output, nullptr);
+        return success();
+      },
+      [](DialectRegistry &registry) {
+        registry.insert<test::TestDialect>();
+        registerBuiltinDialectTranslation(registry);
+        registerLLVMDialectTranslation(registry);
+        registry.addExtension(
+            +[](MLIRContext *ctx, test::TestDialect *dialect) {
+              dialect->addInterfaces<TestDialectLLVMIRTranslationInterface>();
+            });
+      });
+}
+
+} // namespace mlir

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

 #include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
 
 using namespace mlir;
 
 namespace mlir {
 // Defined in the test directory, no public header.
 void registerTestRoundtripSPIRV();
 void registerTestRoundtripDebugSPIRV();
+#ifdef MLIR_INCLUDE_TESTS
+void registerTestToLLVMIR();
+#endif
 } // namespace mlir
 
 static void registerTestTranslations() {
   registerTestRoundtripSPIRV();
   registerTestRoundtripDebugSPIRV();
+#ifdef MLIR_INCLUDE_TESTS
+  registerTestToLLVMIR();
+#endif
 }
 
 int main(int argc, char **argv) {
   registerAllTranslations();
   registerTestTranslations();
   return failed(mlirTranslateMain(argc, argv, "MLIR Translation Testing Tool"));
 }