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

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. Testing
infrastructure to allow translating the Test dialect to LLVM IR is added, so
that it can be checked that the current behavior is not broken in the future.

Differential Revision: https://reviews.llvm.org/D158278

5 files changed, 174 insertions, 6 deletions

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 8b38cf7..ce9bf2f 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1377,6 +1377,15 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
   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()))
@@ -1387,7 +1396,6 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
     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) &&
@@ -1403,10 +1411,6 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
   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;
 
diff --git a/mlir/test/Target/LLVMIR/test.mlir b/mlir/test/Target/LLVMIR/test.mlir
new file mode 100644
index 0000000..f48738f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/test.mlir
@@ -0,0 +1,18 @@
+// 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
diff --git a/mlir/test/lib/Dialect/Test/CMakeLists.txt b/mlir/test/lib/Dialect/Test/CMakeLists.txt
index a185cc3..e4cbbda 100644
--- a/mlir/test/lib/Dialect/Test/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Test/CMakeLists.txt
@@ -2,6 +2,7 @@ set(LLVM_OPTIONAL_SOURCES
   TestDialect.cpp
   TestPatterns.cpp
   TestTraits.cpp
+  TestToLLVMIRTranslation.cpp
 )
 
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
@@ -83,4 +84,19 @@ add_mlir_library(MLIRTestDialect
   MLIRTensorDialect
   MLIRTransformUtils
   MLIRTransforms
-)
+  )
+
+add_mlir_translation_library(MLIRTestToLLVMIRTranslation
+  TestToLLVMIRTranslation.cpp
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLLVMDialect
+  MLIRTestDialect
+  MLIRSupport
+  MLIRBuiltinToLLVMIRTranslation
+  MLIRLLVMToLLVMIRTranslation
+  )
diff --git a/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
new file mode 100644
index 0000000..7110d999
--- /dev/null
+++ b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
@@ -0,0 +1,124 @@
+//===- 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
diff --git a/mlir/tools/mlir-translate/mlir-translate.cpp b/mlir/tools/mlir-translate/mlir-translate.cpp
index f39bf58..4f9661c 100644
--- a/mlir/tools/mlir-translate/mlir-translate.cpp
+++ b/mlir/tools/mlir-translate/mlir-translate.cpp
@@ -21,11 +21,17 @@ 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) {