[OffloadPackager] Add option to extract files from images

[jhuber6]

Summary

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.

[jhuber6]

llvm/tools/llvm-objcopy/ObjcopyOpts.td

200 ↗

(On Diff #443703)

I think the syntax for this flag could be improved, users may want to dump all files matching a triple, architecture, etc. Let me know if you have any suggestions.

[jhuber6]

Leftover objdump in test

[tra]

It would be great to allow dumping all offloading binaries at once. E.g. something like --dump-offloading all=file-name or just --dump-offloading file-name which would save all binaries into file-name.0 .. file-name.N.

llvm/tools/llvm-objcopy/ObjcopyOptions.cpp

865 ↗

(On Diff #443705)

--dump-section -> --dump-offloading ?

[MaskRay]

The feature doesn't fit into llvm-objcopy's purpose. llvm-objcopy is a binary manipulation tool which only understands ELF. --dump-offloading seems to parse the content using knowledge of the underlying format.

[jhuber6]

In D129507#3643246, @MaskRay wrote:

The feature doesn't fit into llvm-objcopy's purpose. llvm-objcopy is a binary manipulation tool which only understands ELF. --dump-offloading seems to parse the content using knowledge of the underlying format.

I see, I simply chose it because it had a similar feature for --dump-sections. Should this be placed in llvm-objdump as well then?

In D129507#3643241, @tra wrote:

It would be great to allow dumping all offloading binaries at once. E.g. something like --dump-offloading all=file-name or just --dump-offloading file-name which would save all binaries into file-name.0 .. file-name.N.

That's a good idea, I'll add it.

[tra]

In D129507#3643246, @MaskRay wrote:

The feature doesn't fit into llvm-objcopy's purpose. llvm-objcopy is a binary manipulation tool which only understands ELF. --dump-offloading seems to parse the content using knowledge of the underlying format.

Which tool would be appropriate for something like this? Do you suggest that we need a new tool?

The "only understands ELF" part may be open to interpretation. If one were to implement an extension to ELF format, would objcopy be the right place to support it? While objdump does have knowledge of the underlying formats, it does not deal with extraction of such data, but rather interprets it for human consumption.
This use case is somewhere in between -- it involves binary data manipulation, which is what objcopy does, but it does require some knowledge of the underlying format.

IMO, the offloading packing format is closer to being an extension of ELF, than being yet another ISA to disassemble (i.e. it's simple, generic and unlikely to change much) and objcopy seems like a decent fit. Offloading data manipulation may eventually be extended to insert new offload binaries, too, which would also fit objcopy functionality nicely.

[jhuber6]

In D129507#3643246, @MaskRay wrote:

The feature doesn't fit into llvm-objcopy's purpose. llvm-objcopy is a binary manipulation tool which only understands ELF. --dump-offloading seems to parse the content using knowledge of the underlying format.

Will this feature be acceptable llvm-objcopy considering we already have code there? I'd like to know where you think this code should live before I put in the effort to make any changes.

[MaskRay]

In D129507#3643307, @tra wrote:

In D129507#3643246, @MaskRay wrote:

The feature doesn't fit into llvm-objcopy's purpose. llvm-objcopy is a binary manipulation tool which only understands ELF. --dump-offloading seems to parse the content using knowledge of the underlying format.

Which tool would be appropriate for something like this? Do you suggest that we need a new tool?

The "only understands ELF" part may be open to interpretation. If one were to implement an extension to ELF format, would objcopy be the right place to support it? While objdump does have knowledge of the underlying formats, it does not deal with extraction of such data, but rather interprets it for human consumption.
This use case is somewhere in between -- it involves binary data manipulation, which is what objcopy does, but it does require some knowledge of the underlying format.

For example, @avl wanted to create a DWARF utility and I vaguely recall there was a discussion whether it applied to llvm-objcopy.
The answer is no as llvm-objcopy shouldn't know how to parse llvm-objcopy.

IMO, the offloading packing format is closer to being an extension of ELF, than being yet another ISA to disassemble (i.e. it's simple, generic and unlikely to change much) and objcopy seems like a decent fit. Offloading data manipulation may eventually be extended to insert new offload binaries, too, which would also fit objcopy functionality nicely.

It's not an extension of ELF. It is a new format with content embedded in an ELF section (according to while (Offset < Buffer->getBufferSize()) { and OffloadBinary::create). I don't really keep up-to-date with the already-plethora of GPU/offloading tools so I do not have an idea out of my head. If adding a new tool seems to be overkill, GPU/offloading folks may need to re-think how to reorganize the tools.
The question isn't raised for the first time why we have so many llvm-*/clang-* tools doing GPU/offloading stuff now.

[MaskRay]

Note: llvm-objcopy --dump-section .text=text a.out rewrites a.out (the output is nearly indistinguishable, but the byte stream may be different). You need llvm-objcopy --dump-section .text=text a.out /dev/null to suppress the rewrite.
This means --dump-offloading likely doesn't have good ergonomics in llvm-objcopy.

[tra]

The question isn't raised for the first time why we have so many llvm-*/clang-* tools doing GPU/offloading stuff now.

Well, it's a bit more complicated, but I do agree that offloading machinery we have is a bit of a mess and could use some streamlining.

That said, the issue at hand here is orthogonal to that. I'll take your word that llvm-objcopy is not the right place for offloading image manipulation.
That leaves some other tool.
The options I see are:

• llvm-objdump -- may not be the best fit as we want to save offload binaries to files, not just print them out. Whether we may ever want to embed offloading binaries is an open question. I think it could be useful, but it's probably a very niche use case, even for those few of us who care what's under the hood in the offloading binaries.
• incorporate functionality into clang-offload-packager.
• write a new tool?

I personally would prefer to avoid creating yet another tool. clang-offload-packager may be a reasonable place for this, considering that it already deals with the new offloading binary packaging style.

[jhuber6]

In D129507#3646755, @MaskRay wrote:

It's not an extension of ELF. It is a new format with content embedded in an ELF section (according to while (Offset < Buffer->getBufferSize()) { and OffloadBinary::create). I don't really keep up-to-date with the already-plethora of GPU/offloading tools so I do not have an idea out of my head. If adding a new tool seems to be overkill, GPU/offloading folks may need to re-think how to reorganize the tools.
The question isn't raised for the first time why we have so many llvm-*/clang-* tools doing GPU/offloading stuff now.

FWIW I've pretty much deprecated three of those GPU/offloading tools and introduced two. We have a lot of these tools mostly because of legacy which I'm hoping will be removed in time. Generally the work I'm doing is noisy, but a net reduction in complexity. The features I'm adding here are not strictly necessary for functionality, but are meant to be convenient tools for users familiar with stuff like cuobjdump from Nvidia. I wouldn't think a new tool is necessary as the actual code required to support these features is relatively small and in-line with many other data formats embedded in an ELF.

In D129507#3646757, @MaskRay wrote:

Note: llvm-objcopy --dump-section .text=text a.out rewrites a.out (the output is nearly indistinguishable, but the byte stream may be different). You need llvm-objcopy --dump-section .text=text a.out /dev/null to suppress the rewrite.
This means --dump-offloading likely doesn't have good ergonomics in llvm-objcopy.

Noted, thanks.

[jhuber6]

In D129507#3646824, @tra wrote:

I personally would prefer to avoid creating yet another tool. clang-offload-packager may be a reasonable place for this, considering that it already deals with the new offloading binary packaging style.

That tool is more or less a straight clone of Nvidia's fatbinary. We could potentially put it there, as it already creates the binaries that we embed inside the host, but it's a little out of scope. I can definitely put it there if people have major objections to putting stuff like this in llvm-objdump I could've made everything work without it, but it made the job of passing this information from the compiler driver to the back-end much more straightforward. Theoretically we could get rid of that and merge the clang-linker-wrapper into lld and bring the number of tools to zero.

[jhenderson]

I'm not familiar with some of the mentioned tools, but I tend to agree with @MaskRay's conclusions.

1. The functionality doesn't really fit in llvm-objcopy because the tool is intended for modifying a file in-place or making a (modified) copy to another place, neither of which is what you'll actually want it for. I think it would be appropriate for extracting the offloading section from an ELF binary as a whole, since that would basically just be --dump-section. Ideally, we'd have the ability to extract by section type, not just section name, but that's getting out-of-scope for this discussion.

2. llvm-objdump shouldn't be writing things to a file (or producing "binary" output that is intended for piping directly to a file). It's goal is printing information about object files. It is appropriate for it to be able to dump the offloading sections in an interpreted format, but not to extract them.

That means the functionality needs adding to a different tool, I feel (whether that tool is a new one or an existing one, I don't know).

[jhuber6]

In D129507#3647420, @jhenderson wrote:

I'm not familiar with some of the mentioned tools, but I tend to agree with @MaskRay's conclusions.

1. The functionality doesn't really fit in llvm-objcopy because the tool is intended for modifying a file in-place or making a (modified) copy to another place, neither of which is what you'll actually want it for. I think it would be appropriate for extracting the offloading section from an ELF binary as a whole, since that would basically just be --dump-section. Ideally, we'd have the ability to extract by section type, not just section name, but that's getting out-of-scope for this discussion.

2. llvm-objdump shouldn't be writing things to a file (or producing "binary" output that is intended for piping directly to a file). It's goal is printing information about object files. It is appropriate for it to be able to dump the offloading sections in an interpreted format, but not to extract them.

That means the functionality needs adding to a different tool, I feel (whether that tool is a new one or an existing one, I don't know).

llvm-objdump has the -s -j <secname> functionality, but that just dumps straight to the screen in hex format. I could put this in the clang-offload-packger, the tools that makes these blobs in the first place, but since it's a clang tool and doesn't care about object files I figured it wouldn't be a good fit. But if people are opposed to this kind of functionality in llvm-objdump, llvm-readobj, or llvm-objcopy I can just put it in there so it's at least available.

[MaskRay]

In D129507#3648113, @jhuber6 wrote:

In D129507#3647420, @jhenderson wrote:

I'm not familiar with some of the mentioned tools, but I tend to agree with @MaskRay's conclusions.

1. The functionality doesn't really fit in llvm-objcopy because the tool is intended for modifying a file in-place or making a (modified) copy to another place, neither of which is what you'll actually want it for. I think it would be appropriate for extracting the offloading section from an ELF binary as a whole, since that would basically just be --dump-section. Ideally, we'd have the ability to extract by section type, not just section name, but that's getting out-of-scope for this discussion.

2. llvm-objdump shouldn't be writing things to a file (or producing "binary" output that is intended for piping directly to a file). It's goal is printing information about object files. It is appropriate for it to be able to dump the offloading sections in an interpreted format, but not to extract them.

That means the functionality needs adding to a different tool, I feel (whether that tool is a new one or an existing one, I don't know).

llvm-objdump has the -s -j <secname> functionality, but that just dumps straight to the screen in hex format. I could put this in the clang-offload-packger, the tools that makes these blobs in the first place, but since it's a clang tool and doesn't care about object files I figured it wouldn't be a good fit. But if people are opposed to this kind of functionality in llvm-objdump, llvm-readobj, or llvm-objcopy I can just put it in there so it's at least available.

I think none of llvm-objdump, llvm-readobj, or llvm-objcopy is applicable. I don't know the offloading ecosystem well, but it seems that if a tool already does something like extracting an offloading section, it seems natural for it to support --dump-offloading if such a debugging functionality is needed.
If clang-offload-packager seems misnomer, since the utility is still new, a rename can be considered. If its main code is in clang/lib, you can add the dump functionality in llvm/lib/Object (to sit beside code which parses/processes some binary formats) and let the tool call that code.

[jhuber6]

In D129507#3652600, @MaskRay wrote:

I think none of llvm-objdump, llvm-readobj, or llvm-objcopy is applicable. I don't know the offloading ecosystem well, but it seems that if a tool already does something like extracting an offloading section, it seems natural for it to support --dump-offloading if such a debugging functionality is needed.
If clang-offload-packager seems misnomer, since the utility is still new, a rename can be considered. If its main code is in clang/lib, you can add the dump functionality in llvm/lib/Object (to sit beside code which parses/processes some binary formats) and let the tool call that code.

Fair enough, I'll just stash this with the clang-offload-packager.

[jhuber6]

Moving to clang packager.

[tra]

clang/test/Driver/offload-packager.c

13–15

So, with no -o option the tool treats --image to both choose the embedded image to dump and the file name to dump it into.
The patch appears to have more dumping modes. E.g. file name appears to be deriveable from the embedded image info. That should be tested, too.

Does user need to specify complete kind/triple/arch? Can I dump a subset of images -- e.g. all images with the same triple, or all with arch starting with sm_7?

It would also be very convenient if we could dump an image by its index in the executable and, also, all images. I suspect these two modes would be the ones used most often interactively.

BTW, we do not seem to have a way to list those embedded images. It would be great to add it.

16

I'd use diff -q or maybe cmp, though I don't know if it's available on Windows. We only want to know if the files differ and do not want it to accidentally dump binary on the screen.

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

78–79

Hooray for c++17!

171

A comment would be useful here. Or conversion of --image=k1=v1,k2=v2,k3=v3... into a kN->vN map could be extracted into a helper function with a meaningful name.

196

Nit: Args.count("file") ? Args["file"] : Saver.save(...) has one less !.

197–199

Is this guaranteed to be unique? I'd add an image index to the name.

207

You could use llvm::copy which works with ranges.

[jhuber6]

Thanks for the comments.

clang/test/Driver/offload-packager.c

13–15

So, with no -o option the tool treats --image to both choose the embedded image to dump and the file name to dump it into.
The patch appears to have more dumping modes. E.g. file name appears to be derivable from the embedded image info. That should be tested, too.

I tried to write a test for this. Given no explicit filename and multiple matches it should dump it as <input>-<triple>-<arch>.o or similar. It works when I test it locally but when I wrote a test for this I don't think it was writing to the correct directory because the lit test couldn't find them.

Does user need to specify complete kind/triple/arch? Can I dump a subset of images -- e.g. all images with the same triple, or all with arch starting with sm_7?

As written now, it should just write every image it finds that matches the image values. So --image=kind=openmp would dump everything whose producer was OpenMP. If it was --image=kind=openmp,file=out.o it would then write to out.o multiple times so you'd just end up with the last one. If you didn't provide a filename it would derive one as above. Not sure if it's worth handling this more intelligently since I think we should keep the user's requested name even if it's ambiguous.

It would also be very convenient if we could dump an image by its index in the executable and, also, all images. I suspect these two modes would be the ones used most often interactively.

BTW, we do not seem to have a way to list those embedded images. It would be great to add it.

That should be was llvm-objdump --offloading does, right? Or do you mean something different from just printing what's embedded.

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

78–79

Indeed.

171

I was thinking about that as it's somewhat obtuse, I'll try to make it clearer.

196

The formatting looked a little weird when I did it that way. Not a huge deal however.

197–199

Not necessarily, you're right that it might be good to use the image's index to provide a unique name in the case of conflicts. As mentioned in another comment I'm not sure if we should modify the output file if the user specified it however.

207

Good point.

[tra]

clang/test/Driver/offload-packager.c

13–15

lit test couldn't find them.

I see. The directory the test is run from is not necessarily the temporary directory. For that matter, such a test will fail in some setups where the source tree is read-only.
I think we need something like ar --output=dir which is used for very similar purposes.

Can I dump a subset of images ?

Given that we can narrow the subset to particular kind/triple/arch, it should do for now.

That should be was llvm-objdump --offloading does, right?

Now that lang-offload-packager got more functionality, it would make sense to make it the tool for listing embedded images, too, IMO.

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

197–199

Agreed, user input takes precedence.

My concern was for the clashes caused by the tool itself. There may conceivably be embedded images with identical properties and we need a way to distinguish between them.

[jhuber6]

clang/test/Driver/offload-packager.c

13–15

I think we need something like ar --output=dir which is used for very similar purposes.

That will work for now, I'll add that and update this patch with better testing.

Now that clang-offload-packager got more functionality, it would make sense to make it the tool for listing embedded images, too, IMO.

It shouldn't be too difficult to add since we'll be iterating through all of these already. I could do that in another patch.

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

197–199

Yeah, I'll add .x to the filename for that reason.

[jhuber6]

Update, still having problems making the test but I figured I'd just update now.

[tra]

In D129507#3730464, @jhuber6 wrote:

Update, still having problems making the test but I figured I'd just update now.

You could try something like this:
https://github.com/llvm/llvm-project/blob/d20e632853ad978e8008747d838d51b51d823a53/lld/test/MachO/stabs.s#L48

# RUN: cd %t && %lld -lSystem test.o foo.o no-debug.o -o test

[jhuber6]

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.

[tra]

clang/test/Driver/offload-packager.c

27

I'd also check that only openmp files got extracted. You may also want to delete all files in the temp directory -- I'm not sure if lit guarantees that it will be clean. I've seen other tests manually cleaning it in some cases.

[jhuber6]

Updating documentation, cleaning up, and adjusting tests.

[bader]

clang/test/Driver/offload-packager.c

2–3

Are nvptx and amdgpu target required for this test?
Latest version of the test invokes clang only for x86 target and clang-offload-packager just adds triple as metadata string without using llvm target. Right?

[jhuber6]

clang/test/Driver/offload-packager.c

2–3

You're right, we could get rid of those. Also this made me realize that the test for the bitcode file isn't actually using bitcode, I should probably fix that.

clang/docs/ClangOffloadPackager.rst

 extend in the future. As mentioned previously, multiple offloading images are
 bundled together by simply concatenating them in this format. Because we have
 the magic bytes and size of each image, we can extract them as-needed.
 
 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
 
   OVERVIEW: A utility for bundling several object files into a single binary.
   The output binary can then be embedded into the host section table
   to create a fatbinary containing offloading code.
   
   USAGE: clang-offload-packager [options]

 =======
 
 This tool simply takes many input files from the ``image`` option and creates a
 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

clang/test/Driver/offload-packager.c

+// 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

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

 
 #include "clang/Basic/Version.h"
 
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #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"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 
 using namespace llvm;
 using namespace llvm::object;
 
 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>."),
-                                       cl::value_desc("file"),
-                                       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>."),
 
 static cl::list<std::string>
     DeviceImages("image",
                  cl::desc("List of key and value arguments. Required keywords "
                           "are 'file' and 'triple'."),
                  cl::value_desc("<key>=<value>,..."),
                  cl::cat(ClangOffloadPackagerCategory));
 
 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 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 (!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;
 }