jhuber6 updated this revision to Diff 454106.
jhuber6 added a comment.
Using @tra's suggestion to use `cd`. I had to make the test not apply to Windows
however, since I had to use `realpath`. But we don't support Windows anyway.
Repository:
rG LLVM Github Monorepo
CHANGES SINCE LAST ACTION
https://reviews.llvm.org/D129507/new/
https://reviews.llvm.org/D129507
Files:
clang/test/Driver/offload-packager.c
clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
Index: clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
===================================================================
--- clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
+++ 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,25 @@
static cl::OptionCategory
ClangOffloadPackagerCategory("clang-offload-packager options");
-static cl::opt<std::string> OutputFile("o", cl::Required,
- cl::desc("Write output to <file>."),
+static cl::opt<std::string> OutputDir("output-dir",
+ cl::desc("Default output to <dir>"),
+ cl::value_desc("dir"),
+ cl::cat(ClangOffloadPackagerCategory));
+
+static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
cl::value_desc("file"),
cl::cat(ClangOffloadPackagerCategory));
+static cl::opt<std::string> InputFile(cl::Positional,
+ cl::desc("Extract from <file>."),
+ cl::value_desc("file"),
+ cl::cat(ClangOffloadPackagerCategory));
+
+static cl::opt<bool> AllowMissing(
+ "allow-missing",
+ cl::desc("Create empty files if images are missing when unpackaging.\n"),
+ cl::init(false), cl::cat(ClangOffloadPackagerCategory));
+
static cl::list<std::string>
DeviceImages("image",
cl::desc("List of key and value arguments. Required keywords "
@@ -49,84 +64,197 @@
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<StringRef, StringRef> getImageArguments(StringRef Image,
+ StringSaver &Saver) {
+ DenseMap<StringRef, StringRef> 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<char, 1024> BinaryData;
raw_svector_ostream OS(BinaryData);
for (StringRef Image : DeviceImages) {
BumpPtrAllocator Alloc;
StringSaver Saver(Alloc);
-
- StringMap<StringRef> 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<StringRef, StringRef> 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<llvm::MemoryBuffer> DeviceImage;
- for (const auto &KeyAndValue : Args) {
- StringRef Key = KeyAndValue.getKey();
+ for (const auto &[Key, Value] : Args) {
if (Key == "file") {
llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> 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<MemoryBuffer> 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<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
FileOutputBuffer::create(OutputFile, BinaryData.size());
if (!OutputOrErr)
- return reportError(OutputOrErr.takeError());
+ return OutputOrErr.takeError();
std::unique_ptr<FileOutputBuffer> 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<SmallVector<std::unique_ptr<OffloadBinary>>>
+extractOffloadFiles(MemoryBufferRef Contents) {
+ if (identify_magic(Contents.getBuffer()) != file_magic::offload_binary)
+ return createStringError(inconvertibleErrorCode(),
+ "Input buffer not an offloading binary");
+ SmallVector<std::unique_ptr<OffloadBinary>> Binaries;
+ uint64_t Offset = 0;
+ // There could be multiple offloading binaries stored at this section.
+ while (Offset < Contents.getBuffer().size()) {
+ std::unique_ptr<MemoryBuffer> 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<std::unique_ptr<MemoryBuffer>> BufferOrErr =
+ MemoryBuffer::getFileOrSTDIN(InputFile);
+ if (std::error_code EC = BufferOrErr.getError())
+ return createFileError(InputFile, EC);
+ std::unique_ptr<MemoryBuffer> 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<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
+ FileOutputBuffer::create(Filename, Binary->getImage().size());
+ if (!OutputOrErr)
+ return OutputOrErr.takeError();
+ std::unique_ptr<FileOutputBuffer> 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 (!OutputDir.empty())
+ sys::fs::set_current_path(!sys::fs::is_directory(OutputDir)
+ ? sys::path::parent_path(OutputDir)
+ : OutputDir);
+
+ 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;
}
Index: clang/test/Driver/offload-packager.c
===================================================================
--- /dev/null
+++ clang/test/Driver/offload-packager.c
@@ -0,0 +1,30 @@
+// 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
+// RUN: diff *-nvptx64-nvidia-cuda-sm_80.1.o %S/Inputs/dummy-elf.o
+// RUN: diff *-packager.c-amdgcn-amd-amdhsa-gfx908.2.o %S/Inputs/dummy-elf.o
+// RUN: diff *-packager.c-amdgcn-amd-amdhsa-gfx90a.3.o %S/Inputs/dummy-elf.o
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
