[clang] bf06295 - [OffloadPackager] Add option to extract files from images

Joseph Huber via cfe-commits cfe-commits at lists.llvm.org
Tue Aug 23 05:57:35 PDT 2022


Author: Joseph Huber
Date: 2022-08-23T12:57:16Z
New Revision: bf062954364d69e0fb8e3e04ef3a65c432e8bd50

URL: https://github.com/llvm/llvm-project/commit/bf062954364d69e0fb8e3e04ef3a65c432e8bd50
DIFF: https://github.com/llvm/llvm-project/commit/bf062954364d69e0fb8e3e04ef3a65c432e8bd50.diff

LOG: [OffloadPackager] Add option to extract files from images

We use the `clang-offload-packager` too bundle many files into a single
binary format containing metadata. This is used for offloading
compilation which may contain multiple device binaries of different
types and architectures in a single file. We use this special binary
format to store these files along with some necessary metadata around
them. We use this format because of the difficulty of determining the
filesize of the various binary inputs that will be passed to the
offloading toolchain rather than engineering a solution for each input.

Previously we only support packaing many files into a single binary.
This patch adds support for doing the reverse by using the same
`--image=` syntax. To unpackage a binary we now present an input file
instead of an output.

Reviewed By: tra

Differential Revision: https://reviews.llvm.org/D129507

Added: 
    clang/test/Driver/offload-packager.c

Modified: 
    clang/docs/ClangOffloadPackager.rst
    clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

Removed: 
    


################################################################################
diff  --git a/clang/docs/ClangOffloadPackager.rst b/clang/docs/ClangOffloadPackager.rst
index 9b17b5d028d09..295eb48ac0dcb 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 0000000000000..c4617d06e93d3
--- /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 8e98fab2a8a4b..0f90263d09345 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<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
 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> 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::list<std::string>
     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<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 
diff erently.
+      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 (!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;
 }


        


More information about the cfe-commits mailing list