diff --git a/clang/test/Driver/offload-packager.c b/clang/test/Driver/offload-packager.c new file mode 100644 --- /dev/null +++ b/clang/test/Driver/offload-packager.c @@ -0,0 +1,16 @@ +// REQUIRES: x86-registered-target +// REQUIRES: nvptx-registered-target +// REQUIRES: amdgpu-registered-target + +// 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 diff --git a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp --- a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp +++ b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp @@ -33,11 +33,20 @@ 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::opt AllowMissing( + "allow-missing", + cl::desc("Create empty files if images are missing when unpackaging.\n"), + cl::init(false), cl::cat(ClangOffloadPackagerCategory)); + static cl::list DeviceImages("image", cl::desc("List of key and value arguments. Required keywords " @@ -49,26 +58,7 @@ 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; - } - - auto reportError = [argv](Error E) { - logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0])); - return EXIT_FAILURE; - }; - +static Error bundleImages() { SmallVector BinaryData; raw_svector_ostream OS(BinaryData); for (StringRef Image : DeviceImages) { @@ -77,18 +67,17 @@ 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); + auto [Key, Value] = Arg.split("="); + if (Args.count(Key)) + Args[Key] = Saver.save(Args[Key] + "," + Value); else - Args[KeyAndValue.first] = KeyAndValue.second; + Args[Key] = Value; } 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; @@ -98,7 +87,7 @@ llvm::ErrorOr> ObjectOrErr = llvm::MemoryBuffer::getFileOrSTDIN(KeyAndValue.getValue()); 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) @@ -115,18 +104,146 @@ } 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); + + 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; + } + + for (const auto &Binary : *BinariesOrErr) { + // 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() + "." + + 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); + std::copy(Binary->getImage().bytes_begin(), + Binary->getImage().bytes_end(), 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; }