[mlir] Convert function signatures before converting globals

Summary: This allows global initializers to reference functions.

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

3 files changed, 24 insertions, 2 deletions

diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index e7223bf7349a..3a701018beb5 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -61,6 +61,8 @@ public:
     LLVM::ensureDistinctSuccessors(m);
 
     T translator(m, std::move(llvmModule));
+    if (failed(translator.convertFunctionSignatures()))
+      return nullptr;
     if (failed(translator.convertGlobals()))
       return nullptr;
     if (failed(translator.convertFunctions()))
@@ -94,6 +96,7 @@ private:
   /// Check whether the module contains only supported ops directly in its body.
   static LogicalResult checkSupportedModuleOps(Operation *m);
 
+  LogicalResult convertFunctionSignatures();
   LogicalResult convertFunctions();
   LogicalResult convertGlobals();
   LogicalResult convertOneFunction(LLVMFuncOp func);
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 08150745e80b..657aa84afe1c 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -783,12 +783,13 @@ LogicalResult ModuleTranslation::checkSupportedModuleOps(Operation *m) {
   return success();
 }
 
-LogicalResult ModuleTranslation::convertFunctions() {
+LogicalResult ModuleTranslation::convertFunctionSignatures() {
   // Lock access to the llvm context.
   llvm::sys::SmartScopedLock<true> scopedLock(
       llvmDialect->getLLVMContextMutex());
+
   // Declare all functions first because there may be function calls that form a
-  // call graph with cycles.
+  // call graph with cycles, or global initializers that reference functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
@@ -802,6 +803,14 @@ LogicalResult ModuleTranslation::convertFunctions() {
       return failure();
   }
 
+  return success();
+}
+
+LogicalResult ModuleTranslation::convertFunctions() {
+  // Lock access to the llvm context.
+  llvm::sys::SmartScopedLock<true> scopedLock(
+      llvmDialect->getLLVMContextMutex());
+
   // Convert functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     // Ignore external functions.
diff --git a/mlir/test/Target/llvmir.mlir b/mlir/test/Target/llvmir.mlir
index ef7349b01377..05b8ef10588a 100644
--- a/mlir/test/Target/llvmir.mlir
+++ b/mlir/test/Target/llvmir.mlir
@@ -1230,3 +1230,13 @@ llvm.func @constant_bf16() -> !llvm<"bfloat"> {
 
 // CHECK: ret bfloat 0xR4120
 
+// -----
+
+llvm.func @address_taken() {
+  llvm.return
+}
+
+llvm.mlir.global internal constant @taker_of_address() : !llvm<"void()*"> {
+  %0 = llvm.mlir.addressof @address_taken : !llvm<"void()*">
+  llvm.return %0 : !llvm<"void()*">
+}