[Clang] Always pass detected CUDA path to 'clang-nvlink-wrapper' (#152789)

Summary:
We always want to use the detected path. The clang driver's detection is
far more sophisticated so we should use that whenever possible. Also
update the usage so we properly fall back to path instead of incorrectly
using the `/bin` if not provided.

Author: jhuber6

clang/lib/Driver/ToolChains/Cuda.cpp

@@ -611,9 +611,12 @@ void NVPTX::Linker::ConstructJob(Compilation &C, const JobAction &JA,
     CmdArgs.push_back(Args.MakeArgString(
         "--pxtas-path=" + Args.getLastArgValue(options::OPT_ptxas_path_EQ)));
 
-  if (Args.hasArg(options::OPT_cuda_path_EQ))
-    CmdArgs.push_back(Args.MakeArgString(
-        "--cuda-path=" + Args.getLastArgValue(options::OPT_cuda_path_EQ)));
+  if (Args.hasArg(options::OPT_cuda_path_EQ) || TC.CudaInstallation.isValid()) {
+    StringRef CudaPath = Args.getLastArgValue(
+        options::OPT_cuda_path_EQ,
+        llvm::sys::path::parent_path(TC.CudaInstallation.getBinPath()));
+    CmdArgs.push_back(Args.MakeArgString("--cuda-path=" + CudaPath));
+  }
 
   // Add paths specified in LIBRARY_PATH environment variable as -L options.
   addDirectoryList(Args, CmdArgs, "-L", "LIBRARY_PATH");

clang/test/Driver/cuda-cross-compiling.c

@@ -104,3 +104,11 @@
 // RUN:   -nogpulib -nogpuinc -### %s 2>&1 | FileCheck -check-prefix=STARTUP %s
 
 // STARTUP: clang-nvlink-wrapper{{.*}}"-lc" "-lm" "{{.*}}crt1.o"
+
+//
+// Test cuda path handling
+//
+// RUN: %clang -target nvptx64-nvidia-cuda -march=sm_52 --cuda-path=%S/Inputs/CUDA/usr/local/cuda \
+// RUN:   -nogpulib -nogpuinc -### %s 2>&1 | FileCheck -check-prefix=PATH %s
+
+// PATH: clang-nvlink-wrapper{{.*}}"--cuda-path={{.*}}/Inputs/CUDA/usr/local/cuda"

clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp

@@ -286,12 +286,16 @@ struct Symbol {
 };
 
 Expected<StringRef> runPTXAs(StringRef File, const ArgList &Args) {
-  std::string CudaPath = Args.getLastArgValue(OPT_cuda_path_EQ).str();
-  std::string GivenPath = Args.getLastArgValue(OPT_ptxas_path_EQ).str();
-  Expected<std::string> PTXAsPath =
-      findProgram(Args, "ptxas", {CudaPath + "/bin", GivenPath});
+  SmallVector<StringRef, 1> SearchPaths;
+  if (Arg *A = Args.getLastArg(OPT_cuda_path_EQ))
+    SearchPaths.push_back(Args.MakeArgString(A->getValue() + Twine("/bin")));
+  if (Arg *A = Args.getLastArg(OPT_ptxas_path_EQ))
+    SearchPaths.push_back(Args.MakeArgString(A->getValue()));
+
+  Expected<std::string> PTXAsPath = findProgram(Args, "ptxas", SearchPaths);
   if (!PTXAsPath)
     return PTXAsPath.takeError();
+
   if (!Args.hasArg(OPT_arch))
     return createStringError(
         "must pass in an explicit nvptx64 gpu architecture to 'ptxas'");
@@ -691,9 +695,11 @@ Error runNVLink(ArrayRef<StringRef> Files, const ArgList &Args) {
   if (Args.hasArg(OPT_lto_emit_asm) || Args.hasArg(OPT_lto_emit_llvm))
     return Error::success();
 
-  std::string CudaPath = Args.getLastArgValue(OPT_cuda_path_EQ).str();
-  Expected<std::string> NVLinkPath =
-      findProgram(Args, "nvlink", {CudaPath + "/bin"});
+  SmallVector<StringRef, 1> SearchPaths;
+  if (Arg *A = Args.getLastArg(OPT_cuda_path_EQ))
+    SearchPaths.push_back(Args.MakeArgString(A->getValue() + Twine("/bin")));
+
+  Expected<std::string> NVLinkPath = findProgram(Args, "nvlink", SearchPaths);
   if (!NVLinkPath)
     return NVLinkPath.takeError();