[OpenMP] Add -Xoffload-linker to forward input to the device linker

We use the clang-linker-wrapper to perform device linking of embedded
offloading object files. This is done by generating those jobs inside of
the linker-wrapper itself. This patch adds an argument in Clang and the
linker-wrapper that allows users to forward input to the device linking
phase. This can either be done for every device linker, or for a
specific target triple. We use the `-Xoffload-linker <arg>` and the
`-Xoffload-linker-<triple> <arg>` syntax to accomplish this.

Reviewed By: markdewing, tra

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

Author: jhuber6

clang/include/clang/Driver/Options.td

@@ -823,6 +823,9 @@ def z : Separate<["-"], "z">, Flags<[LinkerInput, RenderAsInput]>,
 def Xlinker : Separate<["-"], "Xlinker">, Flags<[LinkerInput, RenderAsInput]>,
   HelpText<"Pass <arg> to the linker">, MetaVarName<"<arg>">,
   Group<Link_Group>;
+def Xoffload_linker : JoinedAndSeparate<["-"], "Xoffload-linker">,
+  HelpText<"Pass <arg> to the offload linkers or the ones idenfied by -<triple>">, 
+  MetaVarName<"<triple> <arg>">, Group<Link_Group>;
 def Xpreprocessor : Separate<["-"], "Xpreprocessor">, Group<Preprocessor_Group>,
   HelpText<"Pass <arg> to the preprocessor">, MetaVarName<"<arg>">;
 def X_Flag : Flag<["-"], "X">, Group<Link_Group>;

clang/lib/Driver/ToolChains/Clang.cpp

@@ -8393,6 +8393,20 @@ void LinkerWrapper::ConstructJob(Compilation &C, const JobAction &JA,
   Linker->ConstructJob(C, JA, Output, Inputs, Args, LinkingOutput);
   const auto &LinkCommand = C.getJobs().getJobs().back();
 
+  // Forward -Xoffload-linker<-triple> arguments to the device link job.
+  for (auto *Arg : Args.filtered(options::OPT_Xoffload_linker)) {
+    StringRef Val = Arg->getValue(0);
+    if (Val.empty())
+      CmdArgs.push_back(
+          Args.MakeArgString(Twine("-device-linker=") + Arg->getValue(1)));
+    else
+      CmdArgs.push_back(Args.MakeArgString(
+          "-device-linker=" +
+          ToolChain::getOpenMPTriple(Val.drop_front()).getTriple() + "=" +
+          Arg->getValue(1)));
+  }
+  Args.ClaimAllArgs(options::OPT_Xoffload_linker);
+
   // Add the linker arguments to be forwarded by the wrapper.
   CmdArgs.push_back("-linker-path");
   CmdArgs.push_back(LinkCommand->getExecutable());

clang/test/Driver/linker-wrapper.c

@@ -80,3 +80,15 @@
 // CUDA: nvlink{{.*}}-m64 -o {{.*}}.out -arch sm_52 {{.*}}.o
 // CUDA: nvlink{{.*}}-m64 -o {{.*}}.out -arch sm_70 {{.*}}.o {{.*}}.o
 // CUDA: fatbinary{{.*}}-64 --create {{.*}}.fatbin --image=profile=sm_52,file={{.*}}.out --image=profile=sm_70,file={{.*}}.out
+
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o \
+// RUN:   -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --dry-run --host-triple x86_64-unknown-linux-gnu -linker-path \
+// RUN:   /usr/bin/ld --device-linker=a --device-linker=nvptx64-nvidia-cuda=b -- \
+// RUN:   %t.o -o a.out 2>&1 | FileCheck %s --check-prefix=LINKER_ARGS
+
+// LINKER_ARGS: lld{{.*}}-flavor gnu --no-undefined -shared -o {{.*}}.out {{.*}}.o a
+// LINKER_ARGS: nvlink{{.*}}-m64 -o {{.*}}.out -arch sm_70 {{.*}}.o a b

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

@@ -104,3 +104,9 @@
 // RUN:     -foffload-lto %s 2>&1 | FileCheck --check-prefix=CHECK-LTO-LIBRARY %s
 
 // CHECK-LTO-LIBRARY: {{.*}}-lomptarget{{.*}}-lomptarget.devicertl
+
+// RUN:   %clang -### --target=x86_64-unknown-linux-gnu -fopenmp --offload-arch=sm_52 -nogpulib \
+// RUN:     -Xoffload-linker a -Xoffload-linker-nvptx64-nvidia-cuda b -Xoffload-linker-nvptx64 c \
+// RUN:     %s 2>&1 | FileCheck --check-prefix=CHECK-XLINKER %s
+
+// CHECK-XLINKER: -device-linker=a{{.*}}-device-linker=nvptx64-nvidia-cuda=b{{.*}}-device-linker=nvptx64-nvidia-cuda=c{{.*}}--

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

@@ -108,6 +108,12 @@ static cl::list<std::string>
               cl::desc("Argument to pass to the ptxas invocation"),
               cl::cat(ClangLinkerWrapperCategory));
 
+static cl::list<std::string>
+    LinkerArgs("device-linker", cl::ZeroOrMore,
+               cl::desc("Arguments to pass to the device linker invocation"),
+               cl::value_desc("<value> or <triple>=<value>"),
+               cl::cat(ClangLinkerWrapperCategory));
+
 static cl::opt<bool> Verbose("v", cl::ZeroOrMore,
                              cl::desc("Verbose output from tools"),
                              cl::init(false),
@@ -226,6 +232,17 @@ void printCommands(ArrayRef<StringRef> CmdArgs) {
     llvm::errs() << *IC << (std::next(IC) != IE ? " " : "\n");
 }
 
+// Forward user requested arguments to the device linking job.
+void renderXLinkerArgs(SmallVectorImpl<StringRef> &Args, StringRef Triple) {
+  for (StringRef Arg : LinkerArgs) {
+    auto TripleAndValue = Arg.split('=');
+    if (TripleAndValue.second.empty())
+      Args.push_back(TripleAndValue.first);
+    else if (TripleAndValue.first == Triple)
+      Args.push_back(TripleAndValue.second);
+  }
+}
+
 std::string getMainExecutable(const char *Name) {
   void *Ptr = (void *)(intptr_t)&getMainExecutable;
   auto COWPath = sys::fs::getMainExecutable(Name, Ptr);
@@ -531,6 +548,7 @@ Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
   for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
+  renderXLinkerArgs(CmdArgs, TheTriple.getTriple());
   if (Error Err = executeCommands(*NvlinkPath, CmdArgs))
     return std::move(Err);
 
@@ -599,6 +617,7 @@ Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
   for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
+  renderXLinkerArgs(CmdArgs, TheTriple.getTriple());
   if (Error Err = executeCommands(*LLDPath, CmdArgs))
     return std::move(Err);
 
@@ -676,6 +695,7 @@ Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
   for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
+  renderXLinkerArgs(CmdArgs, TheTriple.getTriple());
   if (Error Err = executeCommands(LinkerUserPath, CmdArgs))
     return std::move(Err);