[OpenMP] Add unbundling of archives containing bundled object files into device specific archives.

[jsjodin]

Summary

This patch adds the file type "a", which together with the --unbundle flag, unbundles an archive of bundled object files into an archive containing the device-specific code.
Example:
clang-offload-bundler --unbundle --inputs=libMyLib.a -type=a -outputs=libMyDeviceLib.a -targets=openmp-amdgcn-amdhsa-gfx906

[jdoerfert]

We need tests for this.

[jsjodin]

Full context diff.

[yaxunl]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

965

Need to add comments about what is in the input file and by what criteria the files are extracted and what format is the output file, does not contain index, etc.

[jsjodin]

Added tests and a comment to the UnbundleArchive function to describe what it does.

[grokos]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

927–931

Does the K_GNU case cover all other OSes? What about Windows?
(we'd like to support Windows as well, like the SYCL branch does)

933–938

Can we add a default case which will return a generic extension (or exit with an "unsupported device" error)? Something like this:

if (Device.contains("gfx"))
  return ".bc";
else if (Device.contains("nvidia-cuda"))
  return ".cubin";

// Maybe also print a warning here
return ".o";

It will help identify things that need to be patched when a new architecture gets support in clang.

[jsjodin]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

927–931

From what I could find out the PE/COFF variant is the same as K_GNU, so it should work.

933–938

Yes, I think that would be fine. I was also considering adding a new flag to specify the file extension instead of having this function. We may want to keep this kind of logic in the driver.

[jsjodin]

Made matching for ".cubin" exension explicit, and created a fallback to ".o" with a warning.

[jsjodin]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

933–938

I opted for now to keep this function and explicitly match for .cubin extension and have a fallback to .o with a warning.

[grokos]

The patch looks good now. I don't have any other comments. We can revisit the driver flag proposal for setting the file extension in the future.

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

939–940

I think a space is missing between archive and members.

[jdoerfert]

Fine from my side. @grokos I guess if you're happy you can accept.

[ABataev]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

Better to use Triple class for parsing triple rather than invent something new.

938–943

No need for else here

1044–1045

What about memory leak here?

[keryell]

Adding Tyker (Gauthier Harnisch) because of SYCL similarity

[jsjodin]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

The use of Triple is pretty confusing in this tool, it is actually more than a triple. It is a pair of a llvm::Triple string and a device e.g. "gfx-906" or "sm_61" separated by a "-". I kept the convention of using Triple.

[jdoerfert]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

I argued that somewhere else before. Are we really sure it makes sense to somehow glue together two things so we don't have to add something new? We have a triple in the module, and if we need a module wide cpu/device identifier, let's add one? I mean, you basically gave the perfect argument of not doing it the way it is implemented: "The use of Triple is pretty confusing in this tool, it is actually more than a triple."
Also, having a global target cpu = "..." in the IR module seems beneficial for other reasons anyway.

[jsjodin]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

Having a target cpu in the module is probably a good idea, but it wouldn't directly solve this unless the llvm::Triple was modified to include the cpu/device. Together with the "kind" it is used to name the ELF sections for the bundling, so it is really 3 things put together that makes up the whole identifier.

[jsjodin]

Fixed string memory leak, replaced if-then-else with conditional expression. Added check in case there is no device code in an archive child, and we can skip processing it.

[jdoerfert]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

Remind me, what is "kind" in this context.
Right, either llvm::Triple needs to hold the cpu/device or we need a new structure on the same level that does. That would match a new top-level module member as well.

[jsjodin]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

It is the offload kind, which is currently "host", "cuda", "hip" or "openmp".

[saiislam]

Ping.

[JonChesterfield]

In D80816#2344616, @saiislam wrote:

Ping.

There's already some outstanding comments - maybe we leave this until you're back from vacation?

[sdmitriev]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

1063

I guess this patch needs to be clang-formatted.

1095–1102

What is the reason for limiting the number of targets and outputs to be extracted to one?

[saiislam]

As @jsjodin is not working on this anymore, I have taken it up in a new patch, D93525 (didn't like the commandeering option).
New clang-formatted patch has a couple of new test cases to elaborate its working further.

All reviewers (and Jan) and subscribers have been copied as such.

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

132–140

To summarize, the target is something like "openmp-amdgcn-amd-amdhsa-gfx906" and getDevice is being used to get the device (gfx906). So, it is a combination of "<kind>-<llvm-target>-<device>".

1063

Done in D93525

1095–1102

For our use case, getting a single device specific archive from this unbundle flag option is sufficient.
Please see new test cases at the end of clang-offload-bundler.c in D93525.

Can you give any use case(s) where we need to generate multiple device specific archives in one run of the tool?

[sdmitriev]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

1095–1102

I assume in general case unbundling for archives will be done in a similar way as for object files and for objects unbundling is currently done into multiple outputs in one run. There are many example of such use cases for object files in clang/test/Driver/openmp-offload.c. So, it seems natural to have the same functionality for archives as well.

[sdmitriev]

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

1095–1102

We can probably reuse code for unbundling archives from https://github.com/intel/llvm branch. This implementation does not have one target/output limitation and it is based on the FileHandler interface as the rest of the file type handlers. I have prepared a draft for llvm.org which is based on the intel/llvm code - https://reviews.llvm.org/D94005

[saiislam]

Thank you, @sdmitriev. I have updated the patch to handle multiple targets/outputs in one run of the tool (https://reviews.llvm.org/D93525?id=314633). Please have a look.

[jsjodin]

This is no longer needed.

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

 #include "clang/Basic/Version.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/Host.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
FilesType("type", cl::Required,
                        "  i   - cpp-output\n"
                        "  ii  - c++-cpp-output\n"
                        "  cui - cuda/hip-output\n"
                        "  d   - dependency\n"
                        "  ll  - llvm\n"
                        "  bc  - llvm-bc\n"
                        "  s   - assembler\n"
                        "  o   - object\n"
+                       "  a   - archive of objects\n"
                        "  gch - precompiled-header\n"
                        "  ast - clang AST file"),
               cl::cat(ClangOffloadBundlerCategory));
 static cl::opt<bool>
     Unbundle("unbundle",
              cl::desc("Unbundle bundled file into several output files.\n"),
              cl::init(false), cl::cat(ClangOffloadBundlerCategory));
 

 }
 static bool hasHostKind(StringRef Target) {
   StringRef OffloadKind;
   StringRef Triple;
   getOffloadKindAndTriple(Target, OffloadKind, Triple);
   return OffloadKind == "host";
 }
 
+static StringRef getTriple(StringRef Target) {
+  StringRef OffloadKind;
+  StringRef Triple;
+  getOffloadKindAndTriple(Target, OffloadKind, Triple);
+  return Triple;
+}
+
+static StringRef getDevice(StringRef Triple) {
+  if (Triple.contains("-")) {
+    auto Split = Triple.rsplit('-');
+    return Split.second;
+  } else {
+    return Triple;
+  }
+}
+
 /// Generic file handler interface.
 class FileHandler {
 public:
   FileHandler() {}
 
   virtual ~FileHandler() {}
 
   /// Update the file handler with information from the header of the bundled
   /// file.
   virtual Error ReadHeader(MemoryBuffer &Input) = 0;
 
   /// Read the marker of the next bundled to be read in the file. The bundle
   /// name is returned if there is one in the file, or `None` if there are no
   /// more bundles to be read.
   virtual Expected<Optional<StringRef>>
   ReadBundleStart(MemoryBuffer &Input) = 0;
 
   /// Read the marker that closes the current bundle.
   virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0;
 
   /// Read the current bundle and write the result into the stream \a OS.
-  virtual Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) = 0;
+  virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
 
   /// Write the header of the bundled file to \a OS based on the information
   /// gathered from \a Inputs.
   virtual Error WriteHeader(raw_fd_ostream &OS,
                             ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) = 0;
 
   /// Write the marker that initiates a bundle for the triple \a TargetTriple to
   /// \a OS.
Expected<Optional<StringRef>> ReadBundleStart(MemoryBuffer &Input) final {
     return CurBundleInfo->first();
   }
 
   Error ReadBundleEnd(MemoryBuffer &Input) final {
     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
     return Error::success();
   }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
     StringRef FC = Input.getBuffer();
     OS.write(FC.data() + CurBundleInfo->second.Offset,
              CurBundleInfo->second.Size);
     return Error::success();
   }
 
   Error WriteHeader(raw_fd_ostream &OS,
while (NextSection != Obj->section_end()) {
       if (*TripleOrErr)
         return **TripleOrErr;
     }
     return None;
   }
 
   Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     Expected<StringRef> ContentOrErr = CurrentSection->getContents();
     if (!ContentOrErr)
       return ContentOrErr.takeError();
     StringRef Content = *ContentOrErr;
 
     // Copy fat object contents to the output when extracting host bundle.
     if (Content.size() == 1u && Content.front() == 0)
       Content = StringRef(Input.getBufferStart(), Input.getBufferSize());
Error ReadBundleEnd(MemoryBuffer &Input) final {
     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars + 1);
     if (TripleEnd != FC.npos)
       // Next time we read after the new line.
       ++ReadChars;
 
     return Error::success();
   }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     StringRef FC = Input.getBuffer();
     size_t BundleStart = ReadChars;
 
     // Find end of the bundle.
     size_t BundleEnd = ReadChars = FC.find(BundleEndString, ReadChars);
 
     StringRef Bundle(&FC.data()[BundleStart], BundleEnd - BundleStart);
     OS << Bundle;
CreateFileHandler(MemoryBuffer &FirstInput) {
   if (FilesType == "ll")
     return std::make_unique<TextFileHandler>(/*Comment=*/";");
   if (FilesType == "bc")
     return std::make_unique<BinaryFileHandler>();
   if (FilesType == "s")
     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
   if (FilesType == "o")
     return CreateObjectFileHandler(FirstInput);
+  if (FilesType == "a")
+    return CreateObjectFileHandler(FirstInput);
   if (FilesType == "gch")
     return std::make_unique<BinaryFileHandler>();
   if (FilesType == "ast")
     return std::make_unique<BinaryFileHandler>();
 
   return createStringError(errc::invalid_argument,
                            "'" + FilesType + "': invalid file type specified");
 }
for (auto &E : Worklist) {
     raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
     if (EC)
       return createFileError(E.second, EC);
   }
 
   return Error::success();
 }
 
+static Archive::Kind getDefaultArchiveKindForHost() {
+  return Triple(sys::getDefaultTargetTriple()).isOSDarwin()
+             ? Archive::K_DARWIN
+             : Archive::K_GNU;
+}
+
+static StringRef getDeviceFileExtension(StringRef Device) {
+  if (Device.contains("gfx"))
+    return ".bc";
+  else
+    return ".cubin";
+}
+
+static StringRef removeExtension(StringRef FileName) {
+  StringRef NoExtFileName;
+  if (FileName.contains("."))
+    NoExtFileName = FileName.rsplit('.').first;
+  else
+    NoExtFileName = FileName;
+
+  return NoExtFileName;
+}
+
+static std::string getDeviceLibraryFileName(StringRef BundleFileName,
+                                            StringRef Device) {
+  StringRef LibName = removeExtension(BundleFileName);
+  StringRef Extension = getDeviceFileExtension(Device);
+
+  std::string Result;
+  Result += LibName;
+  Result += Extension;
+  return Result;
+}
+
+static bool checkDeviceOptions(StringRef Device, std::string OffloadArch) {
+  return !OffloadArch.empty() && OffloadArch == Device;
+}
+
+static Error UnbundleArchive() {
+  std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
+  std::string OffloadArch = getDevice(TargetNames.front()).str();
+  std::vector<NewArchiveMember> ArchiveMembers;
+
+  StringRef IFName = InputFileNames.front();
+  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
+      MemoryBuffer::getFileOrSTDIN(IFName, -1, false);
+  if (std::error_code EC = BufOrErr.getError())
+    return createFileError(InputFileNames.front(), EC);
+
+  ArchiveBuffers.push_back(std::move(*BufOrErr));
+  auto LibOrErr =
+      Archive::create(ArchiveBuffers.back()->getMemBufferRef());
+  if (!LibOrErr)
+    return LibOrErr.takeError();
+
+  auto Archive = std::move(*LibOrErr);
+
+  Error ArchiveErr = Error::success();
+  auto ChildEnd = Archive->child_end();
+  for (auto ChildIter = Archive->child_begin(ArchiveErr);
+       ChildIter != ChildEnd; ++ChildIter) {
+    if (ArchiveErr)
+      return ArchiveErr;
+    auto ChildNameOrErr = (*ChildIter).getName();
+    if (!ChildNameOrErr)
+      return ChildNameOrErr.takeError();
+
+    StringRef ChildName = sys::path::filename(*ChildNameOrErr);
+
+    auto ChildBufferRefOrErr = (*ChildIter).getMemoryBufferRef();
+    if (!ChildBufferRefOrErr)
+      return ChildBufferRefOrErr.takeError();
+
+    auto ChildBuffer =
+        MemoryBuffer::getMemBuffer(*ChildBufferRefOrErr, false);
+
+    Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
+      CreateFileHandler(*ChildBuffer);
+    if (!FileHandlerOrErr)
+      return FileHandlerOrErr.takeError();
+
+    std::unique_ptr<FileHandler> &FileHandler = *FileHandlerOrErr;
+    assert(FileHandler);
+
+    if (Error ReadErr = FileHandler.get()->ReadHeader(*ChildBuffer))
+      return ReadErr;
+
+    Expected<Optional<StringRef>> CurTripleOrErr =
+      FileHandler->ReadBundleStart(*ChildBuffer);
+    if (!CurTripleOrErr)
+      return  CurTripleOrErr.takeError();
+
+    StringRef CurKindTriple = **CurTripleOrErr;
+    assert(!CurKindTriple.empty());
+
+    while (!CurKindTriple.empty()) {
+      if (hasHostKind(CurKindTriple)) {
+        // Do nothing, we don't extract host code yet
+      } else if (checkDeviceOptions(getDevice(getTriple(CurKindTriple)),
+                                    OffloadArch)) {
+        std::string BundleData;
+        raw_string_ostream DataStream(BundleData);
+        if (Error Err = FileHandler.get()->ReadBundle(DataStream, *ChildBuffer))
+          return Err;
+
+        std::string *LibraryName =
+            new std::string(getDeviceLibraryFileName(ChildName, OffloadArch));
+        auto MemBuf =
+            MemoryBuffer::getMemBufferCopy(DataStream.str(), *LibraryName);
+        ArchiveBuffers.push_back(std::move(MemBuf));
+        auto MemBufRef = MemoryBufferRef(*(ArchiveBuffers.back()));
+        ArchiveMembers.push_back(NewArchiveMember(MemBufRef));
+      }
+      if (Error Err = FileHandler.get()->ReadBundleEnd(*ChildBuffer))
+        return Err;
+
+      Expected<Optional<StringRef>> NextTripleOrErr =
+        FileHandler->ReadBundleStart(*ChildBuffer);
+      if (!NextTripleOrErr)
+        return NextTripleOrErr.takeError();
+
+      CurKindTriple = ((*NextTripleOrErr).hasValue()) ? **NextTripleOrErr : "";
+    }
+  }
+  assert(!ArchiveErr);
+
+    std::string FileName = OutputFileNames.front();
+  if (Error WriteErr =
+      writeArchive(FileName, ArchiveMembers, true,
+                   getDefaultArchiveKindForHost(), true, false, nullptr))
+    return WriteErr;
+
+  return Error::success();
+}
+
 static void PrintVersion(raw_ostream &OS) {
   OS << clang::getClangToolFullVersion("clang-offload-bundler") << '\n';
 }
 
 int main(int argc, const char **argv) {
   sys::PrintStackTraceOnErrorSignal(argv[0]);
 
   cl::HideUnrelatedOptions(ClangOffloadBundlerCategory);
   cl::SetVersionPrinter(PrintVersion);
   cl::ParseCommandLineOptions(
int main(int argc, const char **argv) {
   bool Error = false;
   if (Unbundle) {
     if (InputFileNames.size() != 1) {
       Error = true;
       reportError(createStringError(
           errc::invalid_argument,
           "only one input file supported in unbundling mode"));
     }
-    if (OutputFileNames.size() != TargetNames.size()) {
+    if (FilesType == "a" && (OutputFileNames.size() != 1 ||
+                             TargetNames.size() != 1)) {
+      Error = true;
+      reportError(createStringError(errc::invalid_argument,
+                                    "number of output files and targets should "
+                                    "be 1 when unbundling an archive"));
+    } else if (OutputFileNames.size() != TargetNames.size()) {
       Error = true;
       reportError(createStringError(errc::invalid_argument,
                                     "number of output files and targets should "
                                     "match in unbundling mode"));
     }
   } else {
     if (OutputFileNames.size() != 1) {
       Error = true;
       reportError(createStringError(
int main(int argc, const char **argv) {
 
   if (Error)
     return 1;
 
   // Save the current executable directory as it will be useful to find other
   // tools.
   BundlerExecutable = sys::fs::getMainExecutable(argv[0], &BundlerExecutable);
 
-  if (llvm::Error Err = Unbundle ? UnbundleFiles() : BundleFiles()) {
+  llvm::Error Err = (Unbundle && FilesType == "a") ?  UnbundleArchive()
+    : (Unbundle) ? UnbundleFiles()
+    : BundleFiles();
+
+  if (Err) {
     reportError(std::move(Err));
     return 1;
   }
   return 0;
 }