[Flang][Driver]Add datalayout before doing LLVM-IR transformation

The earlier available datalyaout allows MLIR to LLVM-IR transformation
to use the datalayout for decisions, such as comparing sizes for
different types of integers.

This should solve https://github.com/llvm/llvm-project/issues/57230

Reviewed By: awarzynski, vzakhari

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

8 files changed, 49 insertions, 13 deletions

diff --git a/flang/include/flang/Frontend/FrontendActions.h b/flang/include/flang/Frontend/FrontendActions.h
index 975aaa0b9da2..3f50d320c1f3 100644
--- a/flang/include/flang/Frontend/FrontendActions.h
+++ b/flang/include/flang/Frontend/FrontendActions.h
@@ -199,7 +199,7 @@ class CodeGenAction : public FrontendAction {
   void executeAction() override;
   /// Runs prescan, parsing, sema and lowers to MLIR.
   bool beginSourceFileAction() override;
-  /// Sets up LLVM's TargetMachine, configures llvmModule accordingly.
+  /// Sets up LLVM's TargetMachine.
   void setUpTargetMachine();
   /// Runs the optimization (aka middle-end) pipeline on the LLVM module
   /// associated with this action.
diff --git a/flang/include/flang/Optimizer/Support/InitFIR.h b/flang/include/flang/Optimizer/Support/InitFIR.h
index bbc50dcec2e7..12bd80e7abd4 100644
--- a/flang/include/flang/Optimizer/Support/InitFIR.h
+++ b/flang/include/flang/Optimizer/Support/InitFIR.h
@@ -32,7 +32,7 @@ namespace fir::support {
       mlir::scf::SCFDialect, mlir::arith::ArithDialect,                        \
       mlir::cf::ControlFlowDialect, mlir::func::FuncDialect,                   \
       mlir::vector::VectorDialect, mlir::math::MathDialect,                    \
-      mlir::complex::ComplexDialect
+      mlir::complex::ComplexDialect, mlir::DLTIDialect
 
 // The definitive list of dialects used by flang.
 #define FLANG_DIALECT_LIST                                                     \
diff --git a/flang/lib/Frontend/CMakeLists.txt b/flang/lib/Frontend/CMakeLists.txt
index 4abca70acaba..fac5f2c1a1f8 100644
--- a/flang/lib/Frontend/CMakeLists.txt
+++ b/flang/lib/Frontend/CMakeLists.txt
@@ -36,6 +36,7 @@ add_flang_library(flangFrontend
   MLIRTransforms
   MLIRLLVMToLLVMIRTranslation
   MLIRSCFToControlFlow
+  MLIRTargetLLVMIRImport
   ${dialect_libs}
 
   LINK_COMPONENTS
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index dfd2089cb9f9..9042332822c8 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -34,6 +34,8 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Target/LLVMIR/Import.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticFrontend.h"
@@ -79,6 +81,16 @@ bool PrescanAndSemaDebugAction::beginSourceFileAction() {
          (generateRtTypeTables() || true);
 }
 
+static void setMLIRDataLayout(mlir::ModuleOp &mlirModule,
+                              const llvm::DataLayout &dl) {
+  mlir::MLIRContext *context = mlirModule.getContext();
+  mlirModule->setAttr(
+      mlir::LLVM::LLVMDialect::getDataLayoutAttrName(),
+      mlir::StringAttr::get(context, dl.getStringRepresentation()));
+  mlir::DataLayoutSpecInterface dlSpec = mlir::translateDataLayout(dl, context);
+  mlirModule->setAttr(mlir::DLTIDialect::kDataLayoutAttrName, dlSpec);
+}
+
 bool CodeGenAction::beginSourceFileAction() {
   llvmCtx = std::make_unique<llvm::LLVMContext>();
   CompilerInstance &ci = this->getInstance();
@@ -123,6 +135,9 @@ bool CodeGenAction::beginSourceFileAction() {
     }
 
     mlirModule = std::make_unique<mlir::ModuleOp>(module.release());
+    setUpTargetMachine();
+    const llvm::DataLayout &dl = tm->createDataLayout();
+    setMLIRDataLayout(*mlirModule, dl);
     return true;
   }
 
@@ -152,10 +167,15 @@ bool CodeGenAction::beginSourceFileAction() {
       kindMap, ci.getInvocation().getLoweringOpts(),
       ci.getInvocation().getFrontendOpts().envDefaults);
 
+  // Fetch module from lb, so we can set
+  mlirModule = std::make_unique<mlir::ModuleOp>(lb.getModule());
+  setUpTargetMachine();
+  const llvm::DataLayout &dl = tm->createDataLayout();
+  setMLIRDataLayout(*mlirModule, dl);
+
   // Create a parse tree and lower it to FIR
   Fortran::parser::Program &parseTree{*ci.getParsing().parseTree()};
   lb.lower(parseTree, ci.getInvocation().getSemanticsContext());
-  mlirModule = std::make_unique<mlir::ModuleOp>(lb.getModule());
 
   // run the default passes.
   mlir::PassManager pm(mlirCtx.get(), mlir::OpPassManager::Nesting::Implicit);
@@ -565,13 +585,7 @@ getCGOptLevel(const Fortran::frontend::CodeGenOptions &opts) {
 void CodeGenAction::setUpTargetMachine() {
   CompilerInstance &ci = this->getInstance();
 
-  // Set the triple based on the CompilerInvocation set-up
   const std::string &theTriple = ci.getInvocation().getTargetOpts().triple;
-  if (llvmModule->getTargetTriple() != theTriple) {
-    ci.getDiagnostics().Report(clang::diag::warn_fe_override_module)
-        << theTriple;
-    llvmModule->setTargetTriple(theTriple);
-  }
 
   // Create `Target`
   std::string error;
@@ -735,6 +749,22 @@ void CodeGenAction::executeAction() {
   if (!llvmModule)
     generateLLVMIR();
 
+  // Set the triple based on the targetmachine (this comes compiler invocation
+  // and the command-line target option if specified, or the default if not
+  // given on the command-line).
+  setUpTargetMachine();
+  const std::string &theTriple = tm->getTargetTriple().str();
+
+  if (llvmModule->getTargetTriple() != theTriple) {
+    ci.getDiagnostics().Report(clang::diag::warn_fe_override_module)
+        << theTriple;
+  }
+  // Always set the triple and data layout, to make sure they match and are set.
+  // Note that this overwrites any datalayout stored in the LLVM-IR. This avoids
+  // an assert for incompatible data layout when the code-generation happens.
+  llvmModule->setTargetTriple(theTriple);
+  llvmModule->setDataLayout(tm->createDataLayout());
+
   // Run LLVM's middle-end (i.e. the optimizer).
   runOptimizationPipeline(*os);
 
@@ -744,9 +774,6 @@ void CodeGenAction::executeAction() {
     return;
   }
 
-  setUpTargetMachine();
-  llvmModule->setDataLayout(tm->createDataLayout());
-
   if (action == BackendActionTy::Backend_EmitBC) {
     // This action has effectively been completed in runOptimizationPipeline.
     return;
diff --git a/flang/test/Driver/emit-llvm.f90 b/flang/test/Driver/emit-llvm.f90
index 8e864421529e..32a5a044f2b0 100644
--- a/flang/test/Driver/emit-llvm.f90
+++ b/flang/test/Driver/emit-llvm.f90
@@ -6,6 +6,7 @@
 ! RUN: %flang_fc1 -emit-llvm %s -o - | FileCheck %s
 
 ! CHECK: ; ModuleID = 'FIRModule'
+! CHECK: target datalayout =
 ! CHECK: define void @_QQmain()
 ! CHECK-NEXT:  ret void
 ! CHECK-NEXT: }
diff --git a/flang/test/Driver/emit-mlir.f90 b/flang/test/Driver/emit-mlir.f90
index 9391195c9433..191ee13396ef 100644
--- a/flang/test/Driver/emit-mlir.f90
+++ b/flang/test/Driver/emit-mlir.f90
@@ -10,6 +10,8 @@
 ! RUN: %flang_fc1 -emit-mlir emit-mlir.f90 && ls emit-mlir.mlir
 
 ! CHECK: module attributes {
+! CHECK-SAME: dlti.dl_spec =
+! CHECK-SAME: llvm.data_layout =
 ! CHECK-LABEL: func @_QQmain() {
 ! CHECK-NEXT:  return
 ! CHECK-NEXT: }
diff --git a/flang/test/Driver/pic-flags.f90 b/flang/test/Driver/pic-flags.f90
index 2f4842f72cad..fb6ab701c382 100644
--- a/flang/test/Driver/pic-flags.f90
+++ b/flang/test/Driver/pic-flags.f90
@@ -1,3 +1,4 @@
+! REQUIRES: aarch64-registered-target && x86-registered-target && arm-registered-target
 ! RUN: %flang -v -S -emit-llvm -o - %s --target=aarch64-linux-gnu -fno-pie 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-STATIC,CHECK-STATIC-IR
 
 ! RUN: %flang -v -S -emit-llvm -o - %s --target=aarch64-linux-gnu 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-PIE-LEVEL2,CHECK-PIE-LEVEL2-IR
@@ -14,7 +15,6 @@
 ! RUN: %flang -v -### -o - %s --target=arm-none-eabi -frwpi 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-RWPI
 ! RUN: %flang -v -### -o - %s --target=arm-none-eabi -fropi -frwpi 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-ROPI-RWPI
 
-
 ! CHECK: -fc1
 
 
diff --git a/flang/unittests/Frontend/FrontendActionTest.cpp b/flang/unittests/Frontend/FrontendActionTest.cpp
index 9b289523e761..f2727656b5eb 100644
--- a/flang/unittests/Frontend/FrontendActionTest.cpp
+++ b/flang/unittests/Frontend/FrontendActionTest.cpp
@@ -178,6 +178,11 @@ TEST_F(FrontendActionTest, EmitLLVM) {
   compInst.getInvocation().getTargetOpts().triple =
       llvm::Triple::normalize(llvm::sys::getDefaultTargetTriple());
 
+  // Initialise LLVM backend
+  llvm::InitializeAllTargets();
+  llvm::InitializeAllTargetMCs();
+  llvm::InitializeAllAsmPrinters();
+
   // Set-up the output stream. We are using output buffer wrapped as an output
   // stream, as opposed to an actual file (or a file descriptor).
   llvm::SmallVector<char> outputFileBuffer;