diff --git a/clang/tools/CMakeLists.txt b/clang/tools/CMakeLists.txt --- a/clang/tools/CMakeLists.txt +++ b/clang/tools/CMakeLists.txt @@ -8,6 +8,7 @@ add_clang_subdirectory(clang-format-vs) add_clang_subdirectory(clang-fuzzer) add_clang_subdirectory(clang-import-test) +add_clang_subdirectory(clang-nvlink-wrapper) add_clang_subdirectory(clang-offload-bundler) add_clang_subdirectory(clang-offload-wrapper) add_clang_subdirectory(clang-scan-deps) diff --git a/clang/tools/clang-nvlink-wrapper/CMakeLists.txt b/clang/tools/clang-nvlink-wrapper/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/clang/tools/clang-nvlink-wrapper/CMakeLists.txt @@ -0,0 +1,25 @@ +set(LLVM_LINK_COMPONENTS BitWriter Core Object Support) + +if(NOT CLANG_BUILT_STANDALONE) + set(tablegen_deps intrinsics_gen) +endif() + +add_clang_executable(clang-nvlink-wrapper + ClangNvlinkWrapper.cpp + + DEPENDS + ${tablegen_deps} + ) + +set(CLANG_NVLINK_WRAPPER_LIB_DEPS + clangBasic + ) + +add_dependencies(clang clang-nvlink-wrapper) + +target_link_libraries(clang-nvlink-wrapper + PRIVATE + ${CLANG_NVLINK_WRAPPER_LIB_DEPS} + ) + +install(TARGETS clang-nvlink-wrapper RUNTIME DESTINATION bin) diff --git a/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp b/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp new file mode 100644 --- /dev/null +++ b/clang/tools/clang-nvlink-wrapper/ClangNvlinkWrapper.cpp @@ -0,0 +1,164 @@ +//===-- 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 heteregenous 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 "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 Help("h", cl::desc("Alias for -help"), cl::Hidden); + +static Error runNVLink(std::string NVLinkPath, + SmallVectorImpl &Args) { + std::vector 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 &Args, + SmallVectorImpl &TmpFiles) { + std::vector> ArchiveBuffers; + + ErrorOr> BufOrErr = + MemoryBuffer::getFileOrSTDIN(Filename, -1, false); + if (std::error_code EC = BufOrErr.getError()) + return createFileError(Filename, EC); + + ArchiveBuffers.push_back(std::move(*BufOrErr)); + Expected> 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 &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(); +} + +int main(int argc, const char **argv) { + sys::PrintStackTraceOnErrorSignal(argv[0]); + + if (Help) { + cl::PrintHelpMessage(); + return 0; + } + + auto reportError = [argv](Error E) { + logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0])); + exit(1); + }; + + ErrorOr NvlinkPath = sys::findProgramByName("nvlink"); + if (!NvlinkPath) { + reportError(createStringError(NvlinkPath.getError(), + "unable to find 'nvlink' in path")); + } + + SmallVector Argv(argv, argv + argc); + SmallVector ArgvSubst; + SmallVector 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]; + 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)) + reportError(std::move(Err)); + if (Error Err = cleanupTmpFiles(TmpFiles)) + reportError(std::move(Err)); + + return 0; +}