[Flang][OpenMP][Driver][MLIR] Port fopenmp-host-ir-file-path flag and add MLIR module attribute to proliferate to OpenMP IR lowering

This patch ports the fopenmp-host-ir-file-path flag from Clang to Flang-new, this flag
is added by the driver to the device pass when doing two phase compilation (device + host).

This flag is then applied to the module when compiling during the OpenMP device phase.
This file can then be utilised during lowering of the OpenMP dialect to LLVM-IR, which
allows the device and host to maintain 1:1 mapping of OpenMP metadata for variables
during lowering via the OpenMPIRBuilders loadOffloadInfoMetadata facilities
(which is used for declare target and I believe target regions as well).

Reviewer: awarzynski

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

6 files changed, 45 insertions, 9 deletions

diff --git a/flang/include/flang/Frontend/LangOptions.h b/flang/include/flang/Frontend/LangOptions.h
index c8edbfd..7adf2ee 100644
--- a/flang/include/flang/Frontend/LangOptions.h
+++ b/flang/include/flang/Frontend/LangOptions.h
@@ -15,6 +15,8 @@
 #ifndef LLVM_FLANG_FRONTEND_LANGOPTIONS_H
 #define LLVM_FLANG_FRONTEND_LANGOPTIONS_H
 
+#include <string>
+
 namespace Fortran::frontend {
 
 /// Bitfields of LangOptions, split out from LangOptions to ensure
@@ -52,6 +54,10 @@ public:
   void set##Name(Type Value) { Name = static_cast<unsigned>(Value); }
 #include "flang/Frontend/LangOptions.def"
 
+  /// Name of the IR file that contains the result of the OpenMP target
+  /// host code generation.
+  std::string OMPHostIRFile;
+
   LangOptions();
 };
 
diff --git a/flang/include/flang/Tools/CrossToolHelpers.h b/flang/include/flang/Tools/CrossToolHelpers.h
index 19d31df..a28b45c 100644
--- a/flang/include/flang/Tools/CrossToolHelpers.h
+++ b/flang/include/flang/Tools/CrossToolHelpers.h
@@ -23,13 +23,14 @@ struct OffloadModuleOpts {
   OffloadModuleOpts() {}
   OffloadModuleOpts(uint32_t OpenMPTargetDebug, bool OpenMPTeamSubscription,
       bool OpenMPThreadSubscription, bool OpenMPNoThreadState,
-      bool OpenMPNoNestedParallelism, bool OpenMPIsDevice)
+      bool OpenMPNoNestedParallelism, bool OpenMPIsDevice,
+      std::string OMPHostIRFile = {})
       : OpenMPTargetDebug(OpenMPTargetDebug),
         OpenMPTeamSubscription(OpenMPTeamSubscription),
         OpenMPThreadSubscription(OpenMPThreadSubscription),
         OpenMPNoThreadState(OpenMPNoThreadState),
         OpenMPNoNestedParallelism(OpenMPNoNestedParallelism),
-        OpenMPIsDevice(OpenMPIsDevice) {}
+        OpenMPIsDevice(OpenMPIsDevice), OMPHostIRFile(OMPHostIRFile) {}
 
   OffloadModuleOpts(Fortran::frontend::LangOptions &Opts)
       : OpenMPTargetDebug(Opts.OpenMPTargetDebug),
@@ -37,7 +38,8 @@ struct OffloadModuleOpts {
         OpenMPThreadSubscription(Opts.OpenMPThreadSubscription),
         OpenMPNoThreadState(Opts.OpenMPNoThreadState),
         OpenMPNoNestedParallelism(Opts.OpenMPNoNestedParallelism),
-        OpenMPIsDevice(Opts.OpenMPIsDevice) {}
+        OpenMPIsDevice(Opts.OpenMPIsDevice), OMPHostIRFile(Opts.OMPHostIRFile) {
+  }
 
   uint32_t OpenMPTargetDebug = 0;
   bool OpenMPTeamSubscription = false;
@@ -45,6 +47,7 @@ struct OffloadModuleOpts {
   bool OpenMPNoThreadState = false;
   bool OpenMPNoNestedParallelism = false;
   bool OpenMPIsDevice = false;
+  std::string OMPHostIRFile = {};
 };
 
 //  Shares assinging of the OpenMP OffloadModuleInterface and its assorted
@@ -59,6 +62,9 @@ void setOffloadModuleInterfaceAttributes(
       offloadMod.setFlags(Opts.OpenMPTargetDebug, Opts.OpenMPTeamSubscription,
           Opts.OpenMPThreadSubscription, Opts.OpenMPNoThreadState,
           Opts.OpenMPNoNestedParallelism);
+
+      if (!Opts.OMPHostIRFile.empty())
+        offloadMod.setHostIRFilePath(Opts.OMPHostIRFile);
     }
   }
 }
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
index c6d5152..6672777 100644
--- a/flang/lib/Frontend/CompilerInvocation.cpp
+++ b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -726,6 +726,16 @@ static bool parseDialectArgs(CompilerInvocation &res, llvm::opt::ArgList &args,
     if (args.hasArg(clang::driver::options::OPT_fopenmp_is_device)) {
       res.getLangOpts().OpenMPIsDevice = 1;
 
+      // Get OpenMP host file path if any and report if a non existent file is
+      // found
+      if (auto *arg = args.getLastArg(
+              clang::driver::options::OPT_fopenmp_host_ir_file_path)) {
+        res.getLangOpts().OMPHostIRFile = arg->getValue();
+        if (!llvm::sys::fs::exists(res.getLangOpts().OMPHostIRFile))
+          diags.Report(clang::diag::err_drv_omp_host_ir_file_not_found)
+              << res.getLangOpts().OMPHostIRFile;
+      }
+
       if (args.hasFlag(
               clang::driver::options::OPT_fopenmp_assume_teams_oversubscription,
               clang::driver::options::
diff --git a/flang/test/Driver/driver-help.f90 b/flang/test/Driver/driver-help.f90
index f57fd32..6551786 100644
--- a/flang/test/Driver/driver-help.f90
+++ b/flang/test/Driver/driver-help.f90
@@ -153,6 +153,8 @@
 ! HELP-FC1-NEXT: -fno-version-loops-for-stride
 ! HELP-FC1-NEXT:                        Do not create unit-strided loops (default)
 ! HELP-FC1-NEXT: -fopenacc              Enable OpenACC
+! HELP-FC1-NEXT: -fopenmp-host-ir-file-path <value> 
+! HELP-FC1-NEXT:                        Path to the IR file produced by the frontend for the host.
 ! HELP-FC1-NEXT: -fopenmp-is-device     Generate code only for an OpenMP target device.
 ! HELP-FC1-NEXT: -fopenmp-target-debug  Enable debugging in the OpenMP offloading device RTL
 ! HELP-FC1-NEXT: -fopenmp               Parse OpenMP pragmas and generate parallel code.
diff --git a/flang/test/Driver/omp-driver-offload.f90 b/flang/test/Driver/omp-driver-offload.f90
index ec10726..c5acd930 100644
--- a/flang/test/Driver/omp-driver-offload.f90
+++ b/flang/test/Driver/omp-driver-offload.f90
@@ -47,15 +47,15 @@
 ! RUN: %flang -### -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa %s 2>&1 | FileCheck --check-prefixes=CHECK-OPENMP-IS-DEVICE %s
 ! CHECK-OPENMP-IS-DEVICE: "{{[^"]*}}flang-new" "-fc1" {{.*}} "-fopenmp" {{.*}} "-fopenmp-is-device" {{.*}}.f90"
 
-! Testing fembed-offload-object and fopenmp-is-device
+! Testing appropriate flags are gnerated and appropriately assigned by the driver when offloading
 ! RUN: %flang -S -### %s -o %t 2>&1 \
 ! RUN: -fopenmp --offload-arch=gfx90a \
 ! RUN: --target=aarch64-unknown-linux-gnu \
-! RUN:   | FileCheck %s --check-prefixes=CHECK-OPENMP-EMBED
-! CHECK-OPENMP-EMBED: "{{[^"]*}}flang-new" "-fc1" "-triple" "aarch64-unknown-linux-gnu" {{.*}} "-fopenmp" {{.*}}.f90"
-! CHECK-OPENMP-EMBED-NEXT: "{{[^"]*}}flang-new" "-fc1" "-triple" "amdgcn-amd-amdhsa" {{.*}} "-fopenmp" {{.*}} "-fopenmp-is-device" {{.*}}.f90"
-! CHECK-OPENMP-EMBED: "{{[^"]*}}clang-offload-packager{{.*}} "--image=file={{.*}}.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp"
-! CHECK-OPENMP-EMBED-NEXT: "{{[^"]*}}flang-new" "-fc1" "-triple" "aarch64-unknown-linux-gnu" {{.*}} "-fopenmp" {{.*}} "-fembed-offload-object={{.*}}.out" {{.*}}.bc"
+! RUN:   | FileCheck %s --check-prefix=OPENMP-OFFLOAD-ARGS
+! OPENMP-OFFLOAD-ARGS: "{{[^"]*}}flang-new" "-fc1" "-triple" "aarch64-unknown-linux-gnu" {{.*}} "-fopenmp" {{.*}}.f90"
+! OPENMP-OFFLOAD-ARGS-NEXT: "{{[^"]*}}flang-new" "-fc1" "-triple" "amdgcn-amd-amdhsa" {{.*}} "-fopenmp" {{.*}} "-fopenmp-host-ir-file-path" "{{.*}}.bc" "-fopenmp-is-device" {{.*}}.f90"
+! OPENMP-OFFLOAD-ARGS: "{{[^"]*}}clang-offload-packager" {{.*}} "--image=file={{.*}}.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp"
+! OPENMP-OFFLOAD-ARGS-NEXT: "{{[^"]*}}flang-new" "-fc1" "-triple" "aarch64-unknown-linux-gnu" {{.*}} "-fopenmp" {{.*}} "-fembed-offload-object={{.*}}.out" {{.*}}.bc"
 
 ! Test -fopenmp with offload for RTL Flag Options
 ! RUN: %flang -### %s -o %t 2>&1 \
@@ -103,3 +103,9 @@
 ! CHECK-RTL-ALL: "{{[^"]*}}flang-new" "-fc1" {{.*}} "-fopenmp" {{.*}} "-fopenmp-is-device" "-fopenmp-target-debug" "-fopenmp-assume-teams-oversubscription"
 ! CHECK-RTL-ALL: "-fopenmp-assume-threads-oversubscription" "-fopenmp-assume-no-thread-state" "-fopenmp-assume-no-nested-parallelism"
 ! CHECK-RTL-ALL: {{.*}}.f90"
+
+! Test diagnostic error when host IR file is non-existent 
+! RUN: not %flang_fc1 %s -o %t 2>&1 -fopenmp -fopenmp-is-device \
+! RUN: -fopenmp-host-ir-file-path non-existant-file.bc \
+! RUN: | FileCheck %s --check-prefix=HOST-IR-MISSING
+! HOST-IR-MISSING: error: provided host compiler IR file 'non-existant-file.bc' is required to generate code for OpenMP target regions but cannot be found
diff --git a/flang/test/Lower/OpenMP/omp-host-ir-flag.f90 b/flang/test/Lower/OpenMP/omp-host-ir-flag.f90
new file mode 100644
index 0000000..f25a024
--- /dev/null
+++ b/flang/test/Lower/OpenMP/omp-host-ir-flag.f90
@@ -0,0 +1,6 @@
+!RUN: %flang_fc1 -emit-llvm-bc -fopenmp -o %t.bc %s 2>&1
+!RUN: %flang_fc1 -emit-mlir -fopenmp -fopenmp-is-device -fopenmp-host-ir-file-path %t.bc -o - %s 2>&1 | FileCheck %s
+
+!CHECK: module attributes {{{.*}}, omp.host_ir_filepath = "{{.*}}.bc", omp.is_device = #omp.isdevice<is_device = true>{{.*}}}
+subroutine omp_subroutine()
+end subroutine omp_subroutine