diff --git a/clang/docs/ClangOffloadPackager.rst b/clang/docs/ClangOffloadPackager.rst index 9b17b5d028d096..295eb48ac0dcbf 100644 --- a/clang/docs/ClangOffloadPackager.rst +++ b/clang/docs/ClangOffloadPackager.rst @@ -145,8 +145,9 @@ Usage ===== This tool can be used with the following arguments. Generally information is -passed as a key-value pair to the ``image=`` argument. The ``file``, ``triple``, -and ``arch`` arguments are considered mandatory to make a valid image. +passed as a key-value pair to the ``image=`` argument. The ``file`` and ``triple``, +arguments are considered mandatory to make a valid image. The ``arch`` argument +is suggested. .. code-block:: console @@ -179,3 +180,12 @@ single output file with all the images combined. .. code-block:: console clang-offload-packager -o out.bin --image=file=input.o,triple=nvptx64,arch=sm_70 + +The inverse operation can be performed instead by passing the packaged binary as +input. In this mode the matching images will either be placed in the output +specified by the ``file`` option. If no ``file`` argument is provided a name +will be generated for each matching image. + +.. code-block:: console + + clang-offload-packager in.bin --image=file=output.o,triple=nvptx64,arch=sm_70 diff --git a/clang/test/Driver/offload-packager.c b/clang/test/Driver/offload-packager.c new file mode 100644 index 00000000000000..c4617d06e93d3d --- /dev/null +++ b/clang/test/Driver/offload-packager.c @@ -0,0 +1,31 @@ +// REQUIRES: x86-registered-target +// REQUIRES: nvptx-registered-target +// REQUIRES: amdgpu-registered-target +// UNSUPPORTED: system-windows + +// Check that we can extract files from the packaged binary. +// RUN: clang-offload-packager -o %t.out \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_80 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx90a \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx90c +// RUN: clang-offload-packager %t.out \ +// RUN: --image=file=%t-sm_70.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \ +// RUN: --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 +// RUN: diff %t-sm_70.o %S/Inputs/dummy-elf.o +// RUN: diff %t-gfx908.o %S/Inputs/dummy-elf.o + +// Check that we generate a new name if one is not given +// RUN: clang-offload-packager -o %t \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_80 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx90a \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=hip,triple=amdgcn-amd-amdhsa,arch=gfx90c +// RUN: cd $(dirname "%t") && clang-offload-packager %t --image=kind=openmp +// RUN: diff *-nvptx64-nvidia-cuda-sm_70.0.o %S/Inputs/dummy-elf.o; rm *-nvptx64-nvidia-cuda-sm_70.0.o +// RUN: diff *-nvptx64-nvidia-cuda-sm_80.1.o %S/Inputs/dummy-elf.o; rm *-nvptx64-nvidia-cuda-sm_80.1.o +// RUN: diff *-amdgcn-amd-amdhsa-gfx908.2.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx908.2.o +// RUN: diff *-amdgcn-amd-amdhsa-gfx90a.3.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx90a.3.o +// RUN: not diff *-amdgcn-amd-amdhsa-gfx90c.4.o %S/Inputs/dummy-elf.o diff --git a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp index 8e98fab2a8a4be..0f90263d09345a 100644 --- a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp +++ b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp @@ -19,6 +19,7 @@ #include "llvm/Object/OffloadBinary.h" #include "llvm/Support/CommandLine.h" #include "llvm/Support/FileOutputBuffer.h" +#include "llvm/Support/FileSystem.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/Path.h" #include "llvm/Support/Signals.h" @@ -33,11 +34,15 @@ static cl::opt Help("h", cl::desc("Alias for -help"), cl::Hidden); static cl::OptionCategory ClangOffloadPackagerCategory("clang-offload-packager options"); -static cl::opt OutputFile("o", cl::Required, - cl::desc("Write output to ."), +static cl::opt OutputFile("o", cl::desc("Write output to ."), cl::value_desc("file"), cl::cat(ClangOffloadPackagerCategory)); +static cl::opt InputFile(cl::Positional, + cl::desc("Extract from ."), + cl::value_desc("file"), + cl::cat(ClangOffloadPackagerCategory)); + static cl::list DeviceImages("image", cl::desc("List of key and value arguments. Required keywords " @@ -49,84 +54,192 @@ static void PrintVersion(raw_ostream &OS) { OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n'; } -int main(int argc, const char **argv) { - sys::PrintStackTraceOnErrorSignal(argv[0]); - cl::HideUnrelatedOptions(ClangOffloadPackagerCategory); - cl::SetVersionPrinter(PrintVersion); - cl::ParseCommandLineOptions( - argc, argv, - "A utility for bundling several object files into a single binary.\n" - "The output binary can then be embedded into the host section table\n" - "to create a fatbinary containing offloading code.\n"); - - if (Help) { - cl::PrintHelpMessage(); - return EXIT_SUCCESS; +// Get a map containing all the arguments for the image. Repeated arguments will +// be placed in a comma separated list. +static DenseMap getImageArguments(StringRef Image, + StringSaver &Saver) { + DenseMap Args; + for (StringRef Arg : llvm::split(Image, ",")) { + auto [Key, Value] = Arg.split("="); + if (Args.count(Key)) + Args[Key] = Saver.save(Args[Key] + "," + Value); + else + Args[Key] = Value; } - auto reportError = [argv](Error E) { - logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0])); - return EXIT_FAILURE; - }; + return Args; +} +static Error bundleImages() { SmallVector BinaryData; raw_svector_ostream OS(BinaryData); for (StringRef Image : DeviceImages) { BumpPtrAllocator Alloc; StringSaver Saver(Alloc); - - StringMap Args; - for (StringRef Arg : llvm::split(Image, ",")) { - auto KeyAndValue = Arg.split("="); - if (Args.count(KeyAndValue.first)) - Args[KeyAndValue.first] = - Saver.save(Args[KeyAndValue.first] + "," + KeyAndValue.second); - else - Args[KeyAndValue.first] = KeyAndValue.second; - } + DenseMap Args = getImageArguments(Image, Saver); if (!Args.count("triple") || !Args.count("file")) - return reportError(createStringError( + return createStringError( inconvertibleErrorCode(), - "'file' and 'triple' are required image arguments")); + "'file' and 'triple' are required image arguments"); OffloadBinary::OffloadingImage ImageBinary{}; std::unique_ptr DeviceImage; - for (const auto &KeyAndValue : Args) { - StringRef Key = KeyAndValue.getKey(); + for (const auto &[Key, Value] : Args) { if (Key == "file") { llvm::ErrorOr> ObjectOrErr = - llvm::MemoryBuffer::getFileOrSTDIN(KeyAndValue.getValue()); + llvm::MemoryBuffer::getFileOrSTDIN(Value); if (std::error_code EC = ObjectOrErr.getError()) - return reportError(errorCodeToError(EC)); + return errorCodeToError(EC); // Clang uses the '.o' suffix for LTO bitcode. if (identify_magic((*ObjectOrErr)->getBuffer()) == file_magic::bitcode) ImageBinary.TheImageKind = object::IMG_Bitcode; else - ImageBinary.TheImageKind = getImageKind( - sys::path::extension(KeyAndValue.getValue()).drop_front()); + ImageBinary.TheImageKind = + getImageKind(sys::path::extension(Value).drop_front()); ImageBinary.Image = std::move(*ObjectOrErr); } else if (Key == "kind") { - ImageBinary.TheOffloadKind = getOffloadKind(KeyAndValue.getValue()); + ImageBinary.TheOffloadKind = getOffloadKind(Value); } else { - ImageBinary.StringData[Key] = KeyAndValue.getValue(); + ImageBinary.StringData[Key] = Value; } } std::unique_ptr Buffer = OffloadBinary::write(ImageBinary); if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0) - return reportError( - createStringError(inconvertibleErrorCode(), - "Offload binary has invalid size alignment")); + return createStringError(inconvertibleErrorCode(), + "Offload binary has invalid size alignment"); OS << Buffer->getBuffer(); } Expected> OutputOrErr = FileOutputBuffer::create(OutputFile, BinaryData.size()); if (!OutputOrErr) - return reportError(OutputOrErr.takeError()); + return OutputOrErr.takeError(); std::unique_ptr Output = std::move(*OutputOrErr); std::copy(BinaryData.begin(), BinaryData.end(), Output->getBufferStart()); if (Error E = Output->commit()) - return reportError(std::move(E)); + return std::move(E); + return Error::success(); +} + +static Expected>> +extractOffloadFiles(MemoryBufferRef Contents) { + if (identify_magic(Contents.getBuffer()) != file_magic::offload_binary) + return createStringError(inconvertibleErrorCode(), + "Input buffer not an offloading binary"); + SmallVector> Binaries; + uint64_t Offset = 0; + // There could be multiple offloading binaries stored at this section. + while (Offset < Contents.getBuffer().size()) { + std::unique_ptr Buffer = + MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "", + /*RequiresNullTerminator*/ false); + auto BinaryOrErr = OffloadBinary::create(*Buffer); + if (!BinaryOrErr) + return BinaryOrErr.takeError(); + + Offset += (*BinaryOrErr)->getSize(); + Binaries.emplace_back(std::move(*BinaryOrErr)); + } + + return Binaries; +} + +static Error unbundleImages() { + ErrorOr> BufferOrErr = + MemoryBuffer::getFileOrSTDIN(InputFile); + if (std::error_code EC = BufferOrErr.getError()) + return createFileError(InputFile, EC); + std::unique_ptr Buffer = std::move(*BufferOrErr); + + // This data can be misaligned if extracted from an archive. + if (!isAddrAligned(Align(OffloadBinary::getAlignment()), + Buffer->getBufferStart())) + Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(), + Buffer->getBufferIdentifier()); + + auto BinariesOrErr = extractOffloadFiles(*Buffer); + if (!BinariesOrErr) + return BinariesOrErr.takeError(); + + // Try to extract each device image specified by the user from the input file. + for (StringRef Image : DeviceImages) { + BumpPtrAllocator Alloc; + StringSaver Saver(Alloc); + auto Args = getImageArguments(Image, Saver); + + for (uint64_t I = 0, E = BinariesOrErr->size(); I != E; ++I) { + const auto &Binary = (*BinariesOrErr)[I]; + // We handle the 'file' and 'kind' identifiers differently. + bool Match = llvm::all_of(Args, [&](auto &Arg) { + const auto [Key, Value] = Arg; + if (Key == "file") + return true; + if (Key == "kind") + return Binary->getOffloadKind() == getOffloadKind(Value); + return Binary->getString(Key) == Value; + }); + if (!Match) + continue; + + // If the user did not provide a filename derive one from the input and + // image. + StringRef Filename = + !Args.count("file") + ? Saver.save(sys::path::stem(InputFile) + "-" + + Binary->getTriple() + "-" + Binary->getArch() + "." + + std::to_string(I) + "." + + getImageKindName(Binary->getImageKind())) + : Args["file"]; + + Expected> OutputOrErr = + FileOutputBuffer::create(Filename, Binary->getImage().size()); + if (!OutputOrErr) + return OutputOrErr.takeError(); + std::unique_ptr Output = std::move(*OutputOrErr); + llvm::copy(Binary->getImage(), Output->getBufferStart()); + if (Error E = Output->commit()) + return std::move(E); + } + } + + return Error::success(); +} + +int main(int argc, const char **argv) { + sys::PrintStackTraceOnErrorSignal(argv[0]); + cl::HideUnrelatedOptions(ClangOffloadPackagerCategory); + cl::SetVersionPrinter(PrintVersion); + cl::ParseCommandLineOptions( + argc, argv, + "A utility for bundling several object files into a single binary.\n" + "The output binary can then be embedded into the host section table\n" + "to create a fatbinary containing offloading code.\n"); + + if (Help) { + cl::PrintHelpMessage(); + return EXIT_SUCCESS; + } + + auto reportError = [argv](Error E) { + logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0])); + return EXIT_FAILURE; + }; + + if (!InputFile.empty() && !OutputFile.empty()) + return reportError( + createStringError(inconvertibleErrorCode(), + "Packaging to an output file and extracting from an " + "input file are mutually exclusive.")); + + if (!OutputFile.empty()) { + if (Error Err = bundleImages()) + return reportError(std::move(Err)); + } else if (!InputFile.empty()) { + if (Error Err = unbundleImages()) + return reportError(std::move(Err)); + } + + return EXIT_SUCCESS; }