| //===-- clang-nvlink-wrapper/ClangNvlinkWrapper.cpp - wrapper over nvlink-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===---------------------------------------------------------------------===// |
| /// |
| /// \file |
| /// This tool works as a wrapper over nvlink program. It transparently passes |
| /// every input option and objects to nvlink except archive files. It reads |
| /// each input archive file to extract archived cubin files as temporary files. |
| /// These temp (*.cubin) files are passed to nvlink, because nvlink does not |
| /// support linking of archive files implicitly. |
| /// |
| /// During linking of heterogeneous device archive libraries, the |
| /// clang-offload-bundler creates a device specific archive of cubin files. |
| /// Such an archive is then passed to this tool to extract cubin files before |
| /// passing to nvlink. |
| /// |
| /// 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-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; |
| NVLArgs.push_back(NVLinkPath); |
| for (auto &Arg : Args) { |
| NVLArgs.push_back(Arg); |
| } |
| |
| if (sys::ExecuteAndWait(NVLinkPath.c_str(), NVLArgs)) |
| return createStringError(inconvertibleErrorCode(), "'nvlink' failed"); |
| return Error::success(); |
| } |
| |
| static Error extractArchiveFiles(StringRef Filename, |
| SmallVectorImpl<std::string> &Args, |
| SmallVectorImpl<std::string> &TmpFiles) { |
| std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers; |
| |
| ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr = |
| MemoryBuffer::getFileOrSTDIN(Filename, false, false); |
| if (std::error_code EC = BufOrErr.getError()) |
| return createFileError(Filename, EC); |
| |
| ArchiveBuffers.push_back(std::move(*BufOrErr)); |
| Expected<std::unique_ptr<llvm::object::Archive>> LibOrErr = |
| object::Archive::create(ArchiveBuffers.back()->getMemBufferRef()); |
| if (!LibOrErr) |
| return LibOrErr.takeError(); |
| |
| auto Archive = std::move(*LibOrErr); |
| |
| Error Err = Error::success(); |
| auto ChildEnd = Archive->child_end(); |
| for (auto ChildIter = Archive->child_begin(Err); ChildIter != ChildEnd; |
| ++ChildIter) { |
| if (Err) |
| return Err; |
| auto ChildNameOrErr = (*ChildIter).getName(); |
| if (!ChildNameOrErr) |
| return ChildNameOrErr.takeError(); |
| |
| StringRef ChildName = sys::path::filename(ChildNameOrErr.get()); |
| |
| auto ChildBufferRefOrErr = (*ChildIter).getMemoryBufferRef(); |
| if (!ChildBufferRefOrErr) |
| return ChildBufferRefOrErr.takeError(); |
| |
| auto ChildBuffer = |
| MemoryBuffer::getMemBuffer(ChildBufferRefOrErr.get(), false); |
| auto ChildNameSplit = ChildName.split('.'); |
| |
| SmallString<16> Path; |
| int FileDesc; |
| if (std::error_code EC = sys::fs::createTemporaryFile( |
| (ChildNameSplit.first), (ChildNameSplit.second), FileDesc, Path)) |
| return createFileError(ChildName, EC); |
| |
| std::string TmpFileName(Path.str()); |
| Args.push_back(TmpFileName); |
| TmpFiles.push_back(TmpFileName); |
| std::error_code EC; |
| raw_fd_ostream OS(Path.c_str(), EC, sys::fs::OF_None); |
| if (EC) |
| return createFileError(TmpFileName, errc::io_error); |
| OS << ChildBuffer->getBuffer(); |
| OS.close(); |
| } |
| return Err; |
| } |
| |
| 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); |
| }; |
| |
| 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 (const std::string &Arg : NVArgs) { |
| if (sys::path::extension(Arg) == ".a") { |
| if (Error Err = extractArchiveFiles(Arg, ArgvSubst, TmpFiles)) |
| reportError(std::move(Err)); |
| } else { |
| ArgvSubst.push_back(Arg); |
| } |
| } |
| |
| 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; |
| } |