[OpenMP] Embed device files into the host IR

[jhuber6]

Summary

This patch adds support for embedding the device object files into the
host IR to create a fat binary. Each offloading file will be inserted
into a section with the following naming format
.llvm.offloading.<triple>.<arch>.<filename>.

Depends on D116542

[jhuber6]

Updating after changing flag in D116542

[jhuber6]

Adding test for multiple input files to embed.

[jhuber6]

Add input filename to the section name to prevent it from being merged if the user does a link.

[jhuber6]

Remove test that was intended for previous commit.

[JonChesterfield]

Description and test have slightly diverged from implementation - filename is appended to disambiguate, but the filecheck regex only looks at the prefix and the name described in the commit message is missing the filename

[jhuber6]

In D116543#3283874, @JonChesterfield wrote:

Description and test have slightly diverged from implementation - filename is appended to disambiguate, but the filecheck regex only looks at the prefix and the name described in the commit message is missing the filename

Thanks, I changed it in the local commit but forgot to copy it over here.

[jhuber6]

Fix test

clang/lib/Driver/ToolChains/Clang.cpp

const char *ModuleName = [&]{
     return ModuleNameArg ? ModuleNameArg->getValue() : "";
   }();
   InputInfo HeaderModuleInput(Inputs[0].getType(), ModuleName, ModuleName);
 
   const InputInfo &Input =
       IsHeaderModulePrecompile ? HeaderModuleInput : Inputs[0];
 
   InputInfoList ModuleHeaderInputs;
+  InputInfoList OpenMPHostInputs;
   const InputInfo *CudaDeviceInput = nullptr;
   const InputInfo *OpenMPDeviceInput = nullptr;
-  const InputInfo *OpenMPHostInput = nullptr;
   for (const InputInfo &I : Inputs) {
     if (&I == &Input) {
       // This is the primary input.
     } else if (IsHeaderModulePrecompile &&
                types::getPrecompiledType(I.getType()) == types::TY_PCH) {
       types::ID Expected = HeaderModuleInput.getType();
       if (I.getType() != Expected) {
         D.Diag(diag::err_drv_module_header_wrong_kind)
             << I.getFilename() << types::getTypeName(I.getType())
             << types::getTypeName(Expected);
       }
       ModuleHeaderInputs.push_back(I);
     } else if ((IsCuda || IsHIP) && !CudaDeviceInput) {
       CudaDeviceInput = &I;
     } else if (IsOpenMPDevice && !OpenMPDeviceInput) {
       OpenMPDeviceInput = &I;
-    } else if (IsOpenMPHost && !OpenMPHostInput) {
-      OpenMPHostInput = &I;
+    } else if (IsOpenMPHost) {
+      OpenMPHostInputs.push_back(I);
     } else {
       llvm_unreachable("unexpectedly given multiple inputs");
     }
   }
 
   const llvm::Triple *AuxTriple =
       (IsCuda || IsHIP) ? TC.getAuxTriple() : nullptr;
   bool IsWindowsMSVC = RawTriple.isWindowsMSVCEnvironment();

   if (IsOpenMPDevice) {
     CmdArgs.push_back("-fopenmp-is-device");
     if (OpenMPDeviceInput) {
       CmdArgs.push_back("-fopenmp-host-ir-file-path");
       CmdArgs.push_back(Args.MakeArgString(OpenMPDeviceInput->getFilename()));
     }
   }
 
+  // Host-side OpenMP offloading recieves the device object files and embeds it
+  // in a named section including the associated target triple and architecture.
+  if (IsOpenMPHost && !OpenMPHostInputs.empty()) {
+    auto InputFile = OpenMPHostInputs.begin();
+    auto OpenMPTCs = C.getOffloadToolChains<Action::OFK_OpenMP>();
+    for (auto TI = OpenMPTCs.first, TE = OpenMPTCs.second; TI != TE;
+         ++TI, ++InputFile) {
+      const ToolChain *TC = TI->second;
+      const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
+      StringRef File =
+          C.getArgs().MakeArgString(TC->getInputFilename(*InputFile));
+      StringRef InputName = Clang::getBaseInputStem(Args, Inputs);
+
+      CmdArgs.push_back(Args.MakeArgString(
+          "-fembed-offload-object=" + File + "," + TC->getTripleString() + "." +
+          TCArgs.getLastArgValue(options::OPT_march_EQ) + "." + InputName));
+    }
+  }
+
   if (Triple.isAMDGPU()) {
     handleAMDGPUCodeObjectVersionOptions(D, Args, CmdArgs);
 
     if (Args.hasFlag(options::OPT_munsafe_fp_atomics,
                      options::OPT_mno_unsafe_fp_atomics, /*Default=*/false))
       CmdArgs.push_back("-munsafe-fp-atomics");
   }
 

clang/test/Driver/openmp-offload-gpu.c

 // RUN:          -fopenmp-new-driver -no-canonical-prefixes -ccc-print-bindings %s -o openmp-offload-gpu 2>&1 \
 // RUN:   | FileCheck -check-prefix=NEW_DRIVER %s
 
 // NEW_DRIVER: "x86_64-unknown-linux-gnu" - "clang", inputs: ["[[HOST_INPUT:.+]]"], output: "[[HOST_BC:.+]]" 
 // NEW_DRIVER: "nvptx64-nvidia-cuda" - "clang", inputs: ["[[DEVICE_INPUT:.+]]", "[[HOST_BC]]"], output: "[[DEVICE_ASM:.+]]"
 // NEW_DRIVER: "nvptx64-nvidia-cuda" - "NVPTX::Assembler", inputs: ["[[DEVICE_ASM]]"], output: "[[DEVICE_OBJ:.+]]" 
 // NEW_DRIVER: "x86_64-unknown-linux-gnu" - "clang", inputs: ["[[HOST_BC]]", "[[DEVICE_OBJ]]"], output: "[[HOST_OBJ:.+]]" 
 // NEW_DRIVER: "x86_64-unknown-linux-gnu" - "[[LINKER:.+]]", inputs: ["[[HOST_OBJ]]"], output: "openmp-offload-gpu"
+
+// RUN:   %clang -### -fopenmp=libomp -fopenmp-targets=nvptx64-nvidia-cuda -Xopenmp-target=nvptx64-nvida-cuda -march=sm_70 \
+// RUN:          --libomptarget-nvptx-bc-path=%S/Inputs/libomptarget/libomptarget-new-nvptx-test.bc \
+// RUN:          -fopenmp-new-driver -no-canonical-prefixes %s -o openmp-offload-gpu 2>&1 \
+// RUN:   | FileCheck -check-prefix=NEW_DRIVER_EMBEDDING %s
+
+// NEW_DRIVER_EMBEDDING: -fembed-offload-object=[[CUBIN:.*\.cubin]],nvptx64-nvidia-cuda.sm_70

clang/test/Frontend/embed-object.ll

 ; RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm \
-; RUN:    -fembed-offload-object=%S/Inputs/empty.h,section -x ir %s -o - \
+; RUN:    -fembed-offload-object=%S/Inputs/empty.h,section1 \
+; RUN:    -fembed-offload-object=%S/Inputs/empty.h,section2 -x ir %s -o - \
 ; RUN:    | FileCheck %s -check-prefix=CHECK
 
-; CHECK: @llvm.embedded.object = private constant [0 x i8] zeroinitializer, section ".llvm.offloading.section"
-; CHECK: @llvm.compiler.used = appending global [2 x i8*] [i8* @x, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @llvm.embedded.object, i32 0, i32 0)], section "llvm.metadata"
+; CHECK: @[[OBJECT1:.+]] = private constant [0 x i8] zeroinitializer, section ".llvm.offloading.section1"
+; CHECK: @[[OBJECT2:.+]] = private constant [0 x i8] zeroinitializer, section ".llvm.offloading.section2"
+; CHECK: @llvm.compiler.used = appending global [3 x i8*] [i8* @x, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @[[OBJECT1]], i32 0, i32 0), i8* getelementptr inbounds ([0 x i8], [0 x i8]* @[[OBJECT2]], i32 0, i32 0)], section "llvm.metadata"
 
 @x = private constant i8 1
 @llvm.compiler.used = appending global [1 x i8*] [i8* @x], section "llvm.metadata"
 
 define i32 @foo() {
   ret i32 0
 }