jhuber6 created this revision.
jhuber6 added reviewers: jdoerfert, tianshilei1992, Meinersbur.
Herald added subscribers: guansong, yaxunl.
jhuber6 requested review of this revision.
Herald added subscribers: cfe-commits, sstefan1.
Herald added a project: clang.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D118944

Files:
  clang/lib/Driver/ToolChains/Clang.cpp
  clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

Index: clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
===================================================================
--- clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -99,8 +99,8 @@
 
 static cl::list<std::string>
     PtxasArgs("ptxas-args", cl::ZeroOrMore,
-                cl::desc("Argument to pass to the ptxas invocation"),
-                cl::cat(ClangLinkerWrapperCategory));
+              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"),
@@ -118,6 +118,10 @@
                                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,
@@ -129,6 +133,9 @@
 /// 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;
 
@@ -507,9 +514,9 @@
 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)
@@ -554,7 +561,10 @@
 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");
@@ -1097,6 +1107,9 @@
     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)
Index: clang/lib/Driver/ToolChains/Clang.cpp
===================================================================
--- clang/lib/Driver/ToolChains/Clang.cpp
+++ clang/lib/Driver/ToolChains/Clang.cpp
@@ -8148,11 +8148,23 @@
                                  const InputInfoList &Inputs,
                                  const ArgList &Args,
                                  const char *LinkingOutput) const {
+  const auto &D = getToolChain().getDriver();
+  const auto TheTriple = getToolChain().getTriple();
+  auto OpenMPTCRange = C.getOffloadToolChains<Action::OFK_OpenMP>();
   ArgStringList CmdArgs;
 
-  if (getToolChain().getDriver().isUsingLTO(/* IsOffload */ true)) {
+  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;
@@ -8165,9 +8177,10 @@
     }
 
     // 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);
 
@@ -8182,7 +8195,7 @@
       BitcodeSuffix += Arch;
 
       ArgStringList BitcodeLibrary;
-      addOpenMPDeviceRTL(D, TCArgs, BitcodeLibrary, BitcodeSuffix,
+      addOpenMPDeviceRTL(TCDriver, TCArgs, BitcodeLibrary, BitcodeSuffix,
                          TC->getTriple());
 
       if (!BitcodeLibrary.empty())
@@ -8210,12 +8223,8 @@
     }
   }
 
-  // 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");
 
@@ -8246,6 +8255,10 @@
   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());
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to