[flang][driver] Add support for -embed-offload-object flag in flang

This patch adds support for the -embed-offload-object flag to embed offloading
binaries in host code. This flag is identical to the clang flag with the same name.

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

Reviewed By: awarzynski, jhuber6

7 files changed, 72 insertions, 0 deletions

diff --git a/flang/include/flang/Frontend/CodeGenOptions.h b/flang/include/flang/Frontend/CodeGenOptions.h
index 7bb10d4..8de150c 100644
--- a/flang/include/flang/Frontend/CodeGenOptions.h
+++ b/flang/include/flang/Frontend/CodeGenOptions.h
@@ -49,6 +49,10 @@ public:
   /// The paths to the pass plugins that were registered using -fpass-plugin.
   std::vector<std::string> LLVMPassPlugins;
 
+  /// List of filenames passed in using the -fembed-offload-object option. These
+  /// are offloading binaries containing device images and metadata.
+  std::vector<std::string> OffloadObjects;
+
   // Define accessors/mutators for code generation options of enumeration type.
 #define CODEGENOPT(Name, Bits, Default)
 #define ENUM_CODEGENOPT(Name, Type, Bits, Default)                             \
diff --git a/flang/include/flang/Frontend/FrontendActions.h b/flang/include/flang/Frontend/FrontendActions.h
index eb9dda7..7a5042b 100644
--- a/flang/include/flang/Frontend/FrontendActions.h
+++ b/flang/include/flang/Frontend/FrontendActions.h
@@ -221,6 +221,9 @@ protected:
   std::unique_ptr<llvm::LLVMContext> llvmCtx;
   std::unique_ptr<llvm::Module> llvmModule;
 
+  /// Embeds offload objects given with specified with -fembed-offload-object
+  void embedOffloadObjects();
+
   /// Generates an LLVM IR module from CodeGenAction::mlirModule and saves it
   /// in CodeGenAction::llvmModule.
   void generateLLVMIR();
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
index 0e9e945..9ab7030 100644
--- a/flang/lib/Frontend/CompilerInvocation.cpp
+++ b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -129,6 +129,11 @@ static void parseCodeGenArgs(Fortran::frontend::CodeGenOptions &opts,
   for (auto *a : args.filtered(clang::driver::options::OPT_fpass_plugin_EQ))
     opts.LLVMPassPlugins.push_back(a->getValue());
 
+  // -fembed-offload-object option
+  for (auto *a :
+       args.filtered(clang::driver::options::OPT_fembed_offload_object_EQ))
+    opts.OffloadObjects.push_back(a->getValue());
+
   // -mrelocation-model option.
   if (const llvm::opt::Arg *A =
           args.getLastArg(clang::driver::options::OPT_mrelocation_model)) {
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index 7e41565..7cb352a 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -47,12 +47,14 @@
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/MC/TargetRegistry.h"
+#include "llvm/Object/OffloadBinary.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <memory>
 
 using namespace Fortran::frontend;
@@ -727,6 +729,25 @@ void CodeGenAction::runOptimizationPipeline(llvm::raw_pwrite_stream &os) {
   mpm.run(*llvmModule, mam);
 }
 
+void CodeGenAction::embedOffloadObjects() {
+  CompilerInstance &ci = this->getInstance();
+  const auto &cgOpts = ci.getInvocation().getCodeGenOpts();
+
+  for (llvm::StringRef offloadObject : cgOpts.OffloadObjects) {
+    llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> objectOrErr =
+        llvm::MemoryBuffer::getFileOrSTDIN(offloadObject);
+    if (std::error_code ec = objectOrErr.getError()) {
+      auto diagID = ci.getDiagnostics().getCustomDiagID(
+          clang::DiagnosticsEngine::Error, "could not open '%0' for embedding");
+      ci.getDiagnostics().Report(diagID) << offloadObject;
+      return;
+    }
+    llvm::embedBufferInModule(
+        *llvmModule, **objectOrErr, ".llvm.offloading",
+        llvm::Align(llvm::object::OffloadBinary::getAlignment()));
+  }
+}
+
 void CodeGenAction::executeAction() {
   CompilerInstance &ci = this->getInstance();
 
@@ -774,12 +795,17 @@ void CodeGenAction::executeAction() {
     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());
 
+  // Embed offload objects specified with -fembed-offload-object
+  if (!ci.getInvocation().getCodeGenOpts().OffloadObjects.empty())
+    embedOffloadObjects();
+
   // Run LLVM's middle-end (i.e. the optimizer).
   runOptimizationPipeline(ci.isOutputStreamNull() ? *os : ci.getOutputStream());
 
diff --git a/flang/test/Driver/driver-help.f90 b/flang/test/Driver/driver-help.f90
index dffc20f..7bdf3e3 100644
--- a/flang/test/Driver/driver-help.f90
+++ b/flang/test/Driver/driver-help.f90
@@ -109,6 +109,8 @@
 ! HELP-FC1-NEXT: -fdefault-double-8  Set the default double precision kind to an 8 byte wide type
 ! HELP-FC1-NEXT: -fdefault-integer-8 Set the default integer kind to an 8 byte wide type
 ! HELP-FC1-NEXT: -fdefault-real-8    Set the default real kind to an 8 byte wide type
+! HELP-FC1-NEXT: -fembed-offload-object=<value>
+! HELP-FC1-NEXT:                        Embed Offloading device-side binary into host object file as a section.
 ! HELP-FC1-NEXT: -ffast-math            Allow aggressive, lossy floating-point optimizations
 ! HELP-FC1-NEXT: -ffixed-form           Process source files in fixed form
 ! HELP-FC1-NEXT: -ffixed-line-length=<value>
diff --git a/flang/test/Driver/embed-error.f90 b/flang/test/Driver/embed-error.f90
new file mode 100644
index 0000000..70690d3
--- /dev/null
+++ b/flang/test/Driver/embed-error.f90
@@ -0,0 +1,12 @@
+
+!----------
+! RUN lines
+!----------
+! Try to embed missing file
+! RUN: not %flang_fc1 -emit-llvm -o - -fembed-offload-object=%S/Inputs/missing.f90 %s 2>&1 | FileCheck %s --check-prefix=ERROR
+
+! ERROR: error: could not open
+
+parameter(i=1)
+integer :: j
+end program
diff --git a/flang/test/Driver/embed.f90 b/flang/test/Driver/embed.f90
new file mode 100644
index 0000000..60d532a
--- /dev/null
+++ b/flang/test/Driver/embed.f90
@@ -0,0 +1,20 @@
+
+!----------
+! RUN lines
+!----------
+! Embed something that can be easily checked
+! RUN: %flang_fc1 -emit-llvm -o - -fembed-offload-object=%S/Inputs/hello.f90 %s 2>&1 | FileCheck %s
+
+! RUN: %flang_fc1 -emit-llvm-bc -o %t.bc %s 2>&1
+! RUN: %flang_fc1 -emit-llvm -o - -fembed-offload-object=%S/Inputs/hello.f90 %t.bc 2>&1 | FileCheck %s
+
+! CHECK: @[[OBJECT_1:.+]] = private constant [61 x i8] c"program hello\0A  write(*,*), \22Hello world!\22\0Aend program hello\0A", section ".llvm.offloading", align 8, !exclude !0
+! CHECK: @llvm.compiler.used = appending global [1 x ptr] [ptr @[[OBJECT_1]]], section "llvm.metadata"
+
+
+! CHECK: !llvm.embedded.objects = !{![[METADATA_1:[0-9]+]]}
+! CHECK: ![[METADATA_1]] = !{ptr @[[OBJECT_1]], !".llvm.offloading"}
+
+parameter(i=1)
+integer :: j
+end program