[clang] f565488 - [Clang][clang-nvlink-wrapper] Pass nvlink path to the wrapper

[Saiyedul Islam via cfe-commits]

Author: Saiyedul Islam
Date: 2021-10-12T16:15:52Z
New Revision: f56548829c4c696d798c252bf097b71538bd45d7

URL: https://github.com/llvm/llvm-project/commit/f56548829c4c696d798c252bf097b71538bd45d7
DIFF: https://github.com/llvm/llvm-project/commit/f56548829c4c696d798c252bf097b71538bd45d7.diff

LOG: [Clang][clang-nvlink-wrapper] Pass nvlink path to the wrapper

Added support of a "--nvlink-path" option in clang-nvlink-wrapper which
takes the path of nvlink binary.

Static Device Library support for OpenMP (D105191) now searches for
nvlink binary and passes its location via this option. In absence
of this option, nvlink binary is searched in locations in PATH.

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

Added: 
    

Modified: 
    clang/lib/Driver/ToolChains/Cuda.cpp
    clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp

Removed: 
    


################################################################################
diff  --git a/clang/lib/Driver/ToolChains/Cuda.cpp b/clang/lib/Driver/ToolChains/Cuda.cpp
index 18351dae39f7e..0ad1ffb079b31 100644
--- a/clang/lib/Driver/ToolChains/Cuda.cpp
+++ b/clang/lib/Driver/ToolChains/Cuda.cpp
@@ -613,6 +613,11 @@ void NVPTX::OpenMPLinker::ConstructJob(Compilation &C, const JobAction &JA,
   AddStaticDeviceLibsLinking(C, *this, JA, Inputs, Args, CmdArgs, "nvptx", GPUArch,
                       false, false);
 
+  // Find nvlink and pass it as "--nvlink-path=" argument of
+  // clang-nvlink-wrapper.
+  CmdArgs.push_back(Args.MakeArgString(
+      Twine("--nvlink-path=" + getToolChain().GetProgramPath("nvlink"))));
+
   const char *Exec =
       Args.MakeArgString(getToolChain().GetProgramPath("clang-nvlink-wrapper"));
   C.addCommand(std::make_unique<Command>(

diff  --git a/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp b/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp
index 5c8b7b9db6884..bc5b9a9f1fde7 100644
--- a/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp
+++ b/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp
@@ -25,6 +25,7 @@
 /// 2. nvlink -o a.out-openmp-nvptx64 /tmp/a.cubin /tmp/b.cubin
 //===---------------------------------------------------------------------===//
 
+#include "clang/Basic/Version.h"
 #include "llvm/Object/Archive.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
@@ -41,6 +42,19 @@ using namespace llvm;
 
 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
 
+// Mark all our options with this category, everything else (except for -help)
+// will be hidden.
+static cl::OptionCategory
+    ClangNvlinkWrapperCategory("clang-nvlink-wrapper options");
+
+static cl::opt<std::string> NvlinkUserPath("nvlink-path",
+                                           cl::desc("Path of nvlink binary"),
+                                           cl::cat(ClangNvlinkWrapperCategory));
+
+// Do not parse nvlink options
+static cl::list<std::string>
+    NVArgs(cl::Sink, cl::desc("<options to be passed to nvlink>..."));
+
 static Error runNVLink(std::string NVLinkPath,
                        SmallVectorImpl<std::string> &Args) {
   std::vector<StringRef> NVLArgs;
@@ -119,8 +133,20 @@ static Error cleanupTmpFiles(SmallVectorImpl<std::string> &TmpFiles) {
   return Error::success();
 }
 
+static void PrintVersion(raw_ostream &OS) {
+  OS << clang::getClangToolFullVersion("clang-nvlink-wrapper") << '\n';
+}
+
 int main(int argc, const char **argv) {
   sys::PrintStackTraceOnErrorSignal(argv[0]);
+  cl::SetVersionPrinter(PrintVersion);
+  cl::HideUnrelatedOptions(ClangNvlinkWrapperCategory);
+  cl::ParseCommandLineOptions(
+      argc, argv,
+      "A wrapper tool over nvlink program. It transparently passes every \n"
+      "input option and objects to nvlink except archive files and path of \n"
+      "nvlink binary. It reads each input archive file to extract archived \n"
+      "cubin files as temporary files.\n");
 
   if (Help) {
     cl::PrintHelpMessage();
@@ -132,12 +158,7 @@ int main(int argc, const char **argv) {
     exit(1);
   };
 
-  ErrorOr<std::string> NvlinkPath = sys::findProgramByName("nvlink");
-  if (!NvlinkPath) {
-    reportError(createStringError(NvlinkPath.getError(),
-                                  "unable to find 'nvlink' in path"));
-  }
-
+  std::string NvlinkPath;
   SmallVector<const char *, 0> Argv(argv, argv + argc);
   SmallVector<std::string, 0> ArgvSubst;
   SmallVector<std::string, 0> TmpFiles;
@@ -145,8 +166,7 @@ int main(int argc, const char **argv) {
   StringSaver Saver(Alloc);
   cl::ExpandResponseFiles(Saver, cl::TokenizeGNUCommandLine, Argv);
 
-  for (size_t i = 1; i < Argv.size(); ++i) {
-    std::string Arg = Argv[i];
+  for (const std::string &Arg : NVArgs) {
     if (sys::path::extension(Arg) == ".a") {
       if (Error Err = extractArchiveFiles(Arg, ArgvSubst, TmpFiles))
         reportError(std::move(Err));
@@ -155,7 +175,19 @@ int main(int argc, const char **argv) {
     }
   }
 
-  if (Error Err = runNVLink(NvlinkPath.get(), ArgvSubst))
+  NvlinkPath = NvlinkUserPath;
+
+  // If user hasn't specified nvlink binary then search it in PATH
+  if (NvlinkPath.empty()) {
+    ErrorOr<std::string> NvlinkPathErr = sys::findProgramByName("nvlink");
+    if (!NvlinkPathErr) {
+      reportError(createStringError(NvlinkPathErr.getError(),
+                                    "unable to find 'nvlink' in path"));
+    }
+    NvlinkPath = NvlinkPathErr.get();
+  }
+
+  if (Error Err = runNVLink(NvlinkPath, ArgvSubst))
     reportError(std::move(Err));
   if (Error Err = cleanupTmpFiles(TmpFiles))
     reportError(std::move(Err));