[mlir] Convert function signatures before converting globals

[silvas]

Summary

This allows global initializers to reference functions.

[ftynse]

Thanks!

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

if (failed(checkSupportedModuleOps(m)))
       return nullptr;
     auto llvmModule = prepareLLVMModule(m);
     if (!llvmModule)
       return nullptr;
 
     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()))
       return nullptr;
 
     return std::move(translator.llvmModule);
   }
 
protected:
 
   /// A helper to look up remapped operands in the value remapping table.
   SmallVector<llvm::Value *, 8> lookupValues(ValueRange values);
 
 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);
   void connectPHINodes(LLVMFuncOp func);
   LogicalResult convertBlock(Block &bb, bool ignoreArguments);
 
   llvm::Constant *getLLVMConstant(llvm::Type *llvmType, Attribute attr,
                                   Location loc);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 
 LogicalResult ModuleTranslation::checkSupportedModuleOps(Operation *m) {
   for (Operation &o : getModuleBody(m).getOperations())
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp>(&o) && !o.isKnownTerminator())
       return o.emitOpError("unsupported module-level operation");
   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(),
         cast<llvm::FunctionType>(function.getType().getUnderlyingType()));
     llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
     functionMapping[function.getName()] = llvmFunc;
 
     // Forward the pass-through attributes to LLVM.
     if (failed(forwardPassthroughAttributes(function.getLoc(),
                                             function.passthrough(), llvmFunc)))
       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.
     if (function.isExternal())
       continue;
 
     if (failed(convertOneFunction(function)))
       return failure();

mlir/test/Target/llvmir.mlir

 // CHECK-LABEL: @constant_bf16
 llvm.func @constant_bf16() -> !llvm<"bfloat"> {
   %0 = llvm.mlir.constant(1.000000e+01 : bf16) : !llvm<"bfloat">
   llvm.return %0 : !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()*">
+}