[OpenMP] Add Cuda path to linker wrapper tool

[jhuber6]

Summary

The linker wrapper tool uses the 'nvlink' and 'ptxas' binaries to link
and assemble device files. Previously we searched for this using the
binaries in the user's path. This didn't work in cases where the user
passed in a specific Cuda path to Clang. This patch changes the linker
wrapper to accept an argument for the Cuda path we can get from Clang.
This should fix #53573.

[Meinersbur]

Thanks for the fix

clang/lib/Driver/ToolChains/Clang.cpp

8151–8152

[style] Avoid "Almost Always Auto"

8161–8162

Since there is no break, would this potentially add multiple --cuda-path?

8184

[nit] unrelated rename

8259–8262

Is moving this relevant?

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

102–103

[nit] unrelated whitespace change

[jhuber6]

clang/lib/Driver/ToolChains/Clang.cpp

8161–8162

Right now the tool chains are initialized based on triples. I don't think it's technically possible for the user to perform an action like -fopenmp-targets=nvptx64-nvidia-cuda,nvptx-nvidia-cuda. The flag is listed ZeroOrMore, so I'm not sure if it's worth explicitly preventing, we'll just take the most recently passed in one.

8184

Avoids shadowing the new D I added.

8259–8262

No, but I wanted to move it so I figured I might as well while I'm here, similar with the other unrelated changes.

[jhuber6]

Address style.

[jhuber6]

Decided to add break.

[tianshilei1992]

clang/lib/Driver/ToolChains/Clang.cpp

8163

Would be better to keep aligned with clang to use --cuda-path.

[jhuber6]

Use long version.

[tianshilei1992]

LG

clang/lib/Driver/ToolChains/Clang.cpp

C.addCommand(std::make_unique<Command>(
       CmdArgs, Inputs, Output));
 }
 
 void LinkerWrapper::ConstructJob(Compilation &C, const JobAction &JA,
                                  const InputInfo &Output,
                                  const InputInfoList &Inputs,
                                  const ArgList &Args,
                                  const char *LinkingOutput) const {
+  const Driver &D = getToolChain().getDriver();
+  const llvm::Triple TheTriple = getToolChain().getTriple();
+  auto OpenMPTCRange = C.getOffloadToolChains<Action::OFK_OpenMP>();
   ArgStringList CmdArgs;
 
-  if (getToolChain().getDriver().isUsingLTO(/* IsOffload */ true)) {
+  // Pass the CUDA path to the linker wrapper tool.
+  for (auto &I : llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
+    const ToolChain *TC = I.second;
+    if (TC->getTriple().isNVPTX()) {
+      CudaInstallationDetector CudaInstallation(D, TheTriple, Args);
+      if (CudaInstallation.isValid())
+        CmdArgs.push_back(Args.MakeArgString(
+            "-cuda-path=" + CudaInstallation.getInstallPath()));
+    }
+  }
+
+  if (D.isUsingLTO(/* IsOffload */ true)) {
     // Pass in target features for each toolchain.
-    auto OpenMPTCRange = C.getOffloadToolChains<Action::OFK_OpenMP>();
     for (auto &I :
          llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
       const ToolChain *TC = I.second;
       const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
       ArgStringList FeatureArgs;
       TC->addClangTargetOptions(TCArgs, FeatureArgs, Action::OFK_OpenMP);
       auto FeatureIt = llvm::find(FeatureArgs, "-target-feature");
       CmdArgs.push_back(Args.MakeArgString(
           "-target-feature=" + TC->getTripleString() + "=" + *(FeatureIt + 1)));
     }
 
     // Pass in the bitcode library to be linked during LTO.
-    for (auto &I : llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
+    for (auto &I :
+         llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
       const ToolChain *TC = I.second;
-      const Driver &D = TC->getDriver();
+      const Driver &TCDriver = TC->getDriver();
       const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
       StringRef Arch = TCArgs.getLastArgValue(options::OPT_march_EQ);
 
       std::string BitcodeSuffix;
       if (TCArgs.hasFlag(options::OPT_fopenmp_target_new_runtime,
                          options::OPT_fno_openmp_target_new_runtime, true))
         BitcodeSuffix += "new-";
       if (TC->getTriple().isNVPTX())
         BitcodeSuffix += "nvptx-";
       else if (TC->getTriple().isAMDGPU())
         BitcodeSuffix += "amdgpu-";
       BitcodeSuffix += Arch;
 
       ArgStringList BitcodeLibrary;
-      addOpenMPDeviceRTL(D, TCArgs, BitcodeLibrary, BitcodeSuffix,
+      addOpenMPDeviceRTL(TCDriver, TCArgs, BitcodeLibrary, BitcodeSuffix,
                          TC->getTriple());
 
       if (!BitcodeLibrary.empty())
         CmdArgs.push_back(
             Args.MakeArgString("-target-library=" + TC->getTripleString() +
                                "-" + Arch + "=" + BitcodeLibrary.back()));
     }
 
if (const Arg *A = Args.getLastArg(options::OPT_O_Group)) {
           OOpt = "2";
       } else if (A->getOption().matches(options::OPT_O0))
         OOpt = "0";
       if (!OOpt.empty())
         CmdArgs.push_back(Args.MakeArgString(Twine("-opt-level=O") + OOpt));
     }
   }
 
-  // Construct the link job so we can wrap around it.
-  Linker->ConstructJob(C, JA, Output, Inputs, Args, LinkingOutput);
-  const auto &LinkCommand = C.getJobs().getJobs().back();
-
   CmdArgs.push_back("-host-triple");
-  CmdArgs.push_back(Args.MakeArgString(getToolChain().getTripleString()));
+  CmdArgs.push_back(Args.MakeArgString(TheTriple.getTriple()));
   if (Args.hasArg(options::OPT_v))
     CmdArgs.push_back("-v");
 
   // Add debug information if present.
   if (const Arg *A = Args.getLastArg(options::OPT_g_Group)) {
     const Option &Opt = A->getOption();
     if (Opt.matches(options::OPT_gN_Group)) {
       if (Opt.matches(options::OPT_gline_directives_only) ||
if (const Arg *A = Args.getLastArg(options::OPT_Rpass_missed_EQ))
     CmdArgs.push_back(
         Args.MakeArgString(Twine("-pass-remarks-missed=") + A->getValue()));
   if (const Arg *A = Args.getLastArg(options::OPT_Rpass_analysis_EQ))
     CmdArgs.push_back(
         Args.MakeArgString(Twine("-pass-remarks-analysis=") + A->getValue()));
   if (Args.getLastArg(options::OPT_save_temps_EQ))
     CmdArgs.push_back("-save-temps");
 
+  // Construct the link job so we can wrap around it.
+  Linker->ConstructJob(C, JA, Output, Inputs, Args, LinkingOutput);
+  const auto &LinkCommand = C.getJobs().getJobs().back();
+
   // Add the linker arguments to be forwarded by the wrapper.
   CmdArgs.push_back("-linker-path");
   CmdArgs.push_back(LinkCommand->getExecutable());
   CmdArgs.push_back("--");
   for (const char *LinkArg : LinkCommand->getArguments())
     CmdArgs.push_back(LinkArg);
 
   const char *Exec =
       Args.MakeArgString(getToolChain().GetProgramPath("clang-linker-wrapper"));
 
   // Replace the executable and arguments of the link job with the
   // wrapper.
   LinkCommand->replaceExecutable(Exec);
   LinkCommand->replaceArguments(CmdArgs);
 }

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

 static cl::opt<std::string>
     HostTriple("host-triple", cl::ZeroOrMore,
                cl::desc("Triple to use for the host compilation"),
                cl::init(sys::getDefaultTargetTriple()),
                cl::cat(ClangLinkerWrapperCategory));
 
 static cl::list<std::string>
     PtxasArgs("ptxas-args", cl::ZeroOrMore,
-                cl::desc("Argument to pass to the ptxas invocation"),
-                cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<bool> Verbose("v", cl::ZeroOrMore,
                              cl::desc("Verbose output from tools"),
                              cl::init(false),
                              cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<DebugKind> DebugInfo(
     cl::desc("Choose debugging level:"), cl::init(NoDebugInfo),
     cl::values(clEnumValN(NoDebugInfo, "g0", "No debug information"),
                clEnumValN(DirectivesOnly, "gline-directives-only",
                           "Direction information"),
                clEnumValN(FullDebugInfo, "g", "Full debugging support")));
 
 static cl::opt<bool> SaveTemps("save-temps", cl::ZeroOrMore,
                                cl::desc("Save intermediary results."),
                                cl::cat(ClangLinkerWrapperCategory));
 
+static cl::opt<std::string> CudaPath("cuda-path", cl::ZeroOrMore,
+                                     cl::desc("Save intermediary results."),
+                                     cl::cat(ClangLinkerWrapperCategory));
+
 // Do not parse linker options.
 static cl::list<std::string>
     HostLinkerArgs(cl::Positional,
                    cl::desc("<options to be passed to linker>..."));
 
 /// Path of the current binary.
 static const char *LinkerExecutable;
 
 /// Filename of the executable being created.
 static StringRef ExecutableName;
 
+/// Binary path for the CUDA installation.
+static std::string CudaBinaryPath;
+
 /// Temporary files created by the linker wrapper.
 static SmallVector<std::string, 16> TempFiles;
 
 /// Codegen flags for LTO backend.
 static codegen::RegisterCodeGenFlags CodeGenFlags;
 
 /// Magic section string that marks the existence of offloading data. The
 /// section string will be formatted as `.llvm.offloading.<triple>.<arch>`.
default:
     return errorCodeToError(object_error::invalid_file_type);
   }
 }
 
 // TODO: Move these to a separate file.
 namespace nvptx {
 Expected<std::string> assemble(StringRef InputFile, Triple TheTriple,
                                StringRef Arch) {
-  // NVPTX uses the nvlink binary to link device object files.
+  // NVPTX uses the ptxas binary to create device object files.
   ErrorOr<std::string> PtxasPath =
-      sys::findProgramByName("ptxas", sys::path::parent_path(LinkerExecutable));
+      sys::findProgramByName("ptxas", {CudaBinaryPath});
   if (!PtxasPath)
     PtxasPath = sys::findProgramByName("ptxas");
   if (!PtxasPath)
     return createStringError(PtxasPath.getError(),
                              "Unable to find 'ptxas' in path");
 
   // Create a new file to write the linked device image to.
   SmallString<128> TempFile;
if (sys::ExecuteAndWait(*PtxasPath, CmdArgs))
     return createStringError(inconvertibleErrorCode(), "'ptxas' failed");
 
   return static_cast<std::string>(TempFile);
 }
 
 Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
                            StringRef Arch) {
   // NVPTX uses the nvlink binary to link device object files.
-  ErrorOr<std::string> NvlinkPath = sys::findProgramByName("nvlink");
+  ErrorOr<std::string> NvlinkPath =
+      sys::findProgramByName("nvlink", {CudaBinaryPath});
+  if (!NvlinkPath)
+    NvlinkPath = sys::findProgramByName("nvlink");
   if (!NvlinkPath)
     return createStringError(NvlinkPath.getError(),
                              "Unable to find 'nvlink' in path");
 
   // Create a new file to write the linked device image to.
   SmallString<128> TempFile;
   if (Error Err =
           createOutputFile(sys::path::filename(ExecutableName) + "-device-" +
if (Help) {
     return EXIT_SUCCESS;
   }
 
   auto reportError = [argv](Error E) {
     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
     return EXIT_FAILURE;
   };
 
+  if (!CudaPath.empty())
+    CudaBinaryPath = CudaPath + "/bin";
+
   ExecutableName = *(llvm::find(HostLinkerArgs, "-o") + 1);
   SmallVector<std::string, 16> LinkerArgs;
   for (const std::string &Arg : HostLinkerArgs)
     LinkerArgs.push_back(Arg);
 
   SmallVector<StringRef, 16> LibraryPaths;
   for (StringRef Arg : LinkerArgs) {
     if (Arg.startswith("-L"))