From f56548829c4c696d798c252bf097b71538bd45d7 Mon Sep 17 00:00:00 2001
From: Saiyedul Islam <Saiyedul.Islam@amd.com>
Date: Sat, 9 Oct 2021 12:06:11 +0000
Subject: [PATCH] [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
---
 clang/lib/Driver/ToolChains/Cuda.cpp          |  5 ++
 .../ClangNvlinkWrapper.cpp                    | 50 +++++++++++++++----
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/clang/lib/Driver/ToolChains/Cuda.cpp b/clang/lib/Driver/ToolChains/Cuda.cpp
index 18351dae39f7..0ad1ffb079b3 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 5c8b7b9db688..bc5b9a9f1fde 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));