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

[saiislam]

Summary

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.

[Meinersbur]

Thanks for the path, but command line parsing should be done properly.

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

51

I suggest a more descriptive flag, path is too general, like --nvlink or --nvlink-command (bugpoint uses --opt-command for the path to opt)

124

[nit] unrelated whitespace change

151

cl::ParseCommandLineOptions(argc, argv,"Program description") must have been called beforehand to make this work.

171

This way of parsing does not use the declared NvlinkUserPath option. Use cl::ParseCommandLineOptions do parse the arguments and

cl::list<std::string> InputFiles(cl::Positional, cl::desc("<input files>..."));

to catch all non-option arguments.

The way it is done here also does not catch the --path <nvlink> (without equal sign) syntax.

[saiislam]

1. Changed the option from path to nvlink-command.
2. Command line arguments are now parsed using proper API.

[saiislam]

Fixed typo

[Meinersbur]

Some nitpicks, but it fixes the build, so LGTM.

clang/lib/Driver/ToolChains/Cuda.cpp

617

[style] LLVM's coding standard does not use almost-always-auto.

It's not immediately obvious here, does GetProgramPath look into the BinPath detected by CudaInstallationDetector? I applied the patch locally to http://meinersbur.de:8011/#/builders/1 and it actually does work.

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

169–170

[tra]

clang/lib/Driver/ToolChains/Cuda.cpp

620

Nit: the variables above are used only once. It would be fine to just push_back(MakeArgString(GetProgramPath())).

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

50

Nit. Clang already has --ptxas-path= option. It may be worth using -path suffix here for consistency, too.

169–170

Nit: for (const std::string &Arg : NVArgs) to avoid unnecessary copies.

[saiislam]

Thanks, Michael and Artem!

1. Changed nvlink-command to nvlink-path.
2. Answered queries and done refactoring as suggested.

[saiislam]

clang-format(ed).

[saiislam]

clang/lib/Driver/ToolChains/Cuda.cpp

617

Yes, you are right. CudaToolChain constructor initializes binary paths obtained from CudaInstallationDetector which is used by this call.

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

50

changed it to --nvlink-path

[Meinersbur]

Thanks, this fixed the errors introduced in D105191: http://meinersbur.de:8011/#/builders/1/builds/1594

clang/lib/Driver/ToolChains/Cuda.cpp

const char *CubinF = C.addTempFile(
         C.getArgs().MakeArgString(getToolChain().getInputFilename(II)));
 
     CmdArgs.push_back(CubinF);
   }
 
   AddStaticDeviceLibsLinking(C, *this, JA, Inputs, Args, CmdArgs, "nvptx", GPUArch,
                       false, false);
 
+  // Find nvlink and pass it as "--nvlink-command=" argument of clang-nvlink-wrapper.
+  auto NvlinkBin = getToolChain().GetProgramPath("nvlink");
+  const char *NvlinkPath =
+      Args.MakeArgString(Twine("--nvlink-command=" + NvlinkBin));
+  CmdArgs.push_back(NvlinkPath);
+
   const char *Exec =
       Args.MakeArgString(getToolChain().GetProgramPath("clang-nvlink-wrapper"));
   C.addCommand(std::make_unique<Command>(
       JA, *this,
       ResponseFileSupport{ResponseFileSupport::RF_Full, llvm::sys::WEM_UTF8,
                           "--options-file"},
       Exec, CmdArgs, Inputs, Output));
 }

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

 ///
 /// Example:
 /// clang-nvlink-wrapper -o a.out-openmp-nvptx64 /tmp/libTest-nvptx-sm_50.a
 ///
 /// 1. Extract (libTest-nvptx-sm_50.a) => /tmp/a.cubin /tmp/b.cubin
 /// 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"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
 
 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-command",
+                                           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;
   NVLArgs.push_back(NVLinkPath);
   for (auto &Arg : Args) {
     NVLArgs.push_back(Arg);
   }
 

 static Error cleanupTmpFiles(SmallVectorImpl<std::string> &TmpFiles) {
   for (auto &TmpFile : TmpFiles) {
     if (std::error_code EC = sys::fs::remove(TmpFile))
       return createFileError(TmpFile, errc::no_such_file_or_directory);
   }
   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();
     return 0;
   }
 
   auto reportError = [argv](Error E) {
     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
     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;
   BumpPtrAllocator Alloc;
   StringSaver Saver(Alloc);
   cl::ExpandResponseFiles(Saver, cl::TokenizeGNUCommandLine, Argv);
 
-  for (size_t i = 1; i < Argv.size(); ++i) {
-    std::string Arg = Argv[i];
+  for (size_t i = 0; i < NVArgs.size(); ++i) {
+    std::string Arg = NVArgs[i];
     if (sys::path::extension(Arg) == ".a") {
       if (Error Err = extractArchiveFiles(Arg, ArgvSubst, TmpFiles))
         reportError(std::move(Err));
     } else {
       ArgvSubst.push_back(Arg);
     }
   }
 
-  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));
 
   return 0;
 }