diff --git a/clang/test/Driver/offload-packager.c b/clang/test/Driver/offload-packager.c --- a/clang/test/Driver/offload-packager.c +++ b/clang/test/Driver/offload-packager.c @@ -53,3 +53,13 @@ // RUN: --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 // RUN: diff %t-sm_70.o %t.elf.o // RUN: diff %t-gfx908.o %t.elf.o + +// Check that we can extract from an archive file to an archive file. +// RUN: clang-offload-packager -o %t.out \ +// RUN: --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \ +// RUN: --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 +// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out +// RUN: llvm-ar rcs %t.a %t.o +// RUN: clang-offload-packager %t.a --archive --image=file=%t-gfx908.a,arch=gfx908 +// RUN: llvm-ar t %t-gfx908.a 2>&1 | FileCheck %s +// CHECK: {{.*}}.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 @@ -15,6 +15,7 @@ #include "clang/Basic/Version.h" #include "llvm/BinaryFormat/Magic.h" +#include "llvm/Object/ArchiveWriter.h" #include "llvm/Object/OffloadBinary.h" #include "llvm/Support/CommandLine.h" #include "llvm/Support/FileOutputBuffer.h" @@ -49,6 +50,14 @@ cl::value_desc("=,..."), cl::cat(ClangOffloadPackagerCategory)); +static cl::opt + CreateArchive("archive", + cl::desc("Write extracted files to a static archive"), + cl::cat(ClangOffloadPackagerCategory)); + +/// Path of the current binary. +static const char *PackagerExecutable; + static void PrintVersion(raw_ostream &OS) { OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n'; } @@ -69,6 +78,18 @@ return Args; } +static Error writeFile(StringRef Filename, StringRef Data) { + Expected> OutputOrErr = + FileOutputBuffer::create(Filename, Data.size()); + if (!OutputOrErr) + return OutputOrErr.takeError(); + std::unique_ptr Output = std::move(*OutputOrErr); + llvm::copy(Data, Output->getBufferStart()); + if (Error E = Output->commit()) + return E; + return Error::success(); +} + static Error bundleImages() { SmallVector BinaryData; raw_svector_ostream OS(BinaryData); @@ -111,13 +132,8 @@ OS << Buffer->getBuffer(); } - Expected> OutputOrErr = - FileOutputBuffer::create(OutputFile, BinaryData.size()); - if (!OutputOrErr) - return OutputOrErr.takeError(); - std::unique_ptr Output = std::move(*OutputOrErr); - std::copy(BinaryData.begin(), BinaryData.end(), Output->getBufferStart()); - if (Error E = Output->commit()) + if (Error E = writeFile(OutputFile, + StringRef(BinaryData.begin(), BinaryData.size()))) return E; return Error::success(); } @@ -145,8 +161,9 @@ StringSaver Saver(Alloc); auto Args = getImageArguments(Image, Saver); - for (uint64_t I = 0, E = Binaries.size(); I != E; ++I) { - const auto *Binary = Binaries[I].getBinary(); + SmallVector Extracted; + for (const OffloadFile &File : Binaries) { + const auto *Binary = File.getBinary(); // We handle the 'file' and 'kind' identifiers differently. bool Match = llvm::all_of(Args, [&](auto &Arg) { const auto [Key, Value] = Arg; @@ -156,27 +173,45 @@ 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()) + if (Match) + Extracted.push_back(Binary); + } + + if (Extracted.empty()) + continue; + + if (CreateArchive) { + if (!Args.count("file")) + return createStringError(inconvertibleErrorCode(), + "Image must have a 'file' argument."); + + SmallVector Members; + for (const OffloadBinary *Binary : Extracted) + Members.emplace_back(MemoryBufferRef( + Binary->getImage(), + Binary->getMemoryBufferRef().getBufferIdentifier())); + + if (Error E = writeArchive(Args["file"], Members, true, + Archive::getDefaultKindForHost(), true, false, + nullptr)) return E; + } else if (Args.count("file")) { + if (Extracted.size() > 1) + WithColor::warning(errs(), PackagerExecutable) + << "Multiple inputs match to a single file, '" << Args["file"] + << "'\n"; + if (Error E = writeFile(Args["file"], Extracted.back()->getImage())) + return E; + } else { + uint64_t Idx = 0; + for (const OffloadBinary *Binary : Extracted) { + StringRef Filename = + Saver.save(sys::path::stem(InputFile) + "-" + Binary->getTriple() + + "-" + Binary->getArch() + "." + std::to_string(Idx++) + + "." + getImageKindName(Binary->getImageKind())); + if (Error E = writeFile(Filename, Binary->getImage())) + return E; + } } } @@ -198,6 +233,7 @@ return EXIT_SUCCESS; } + PackagerExecutable = argv[0]; auto reportError = [argv](Error E) { logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0])); return EXIT_FAILURE;