saiislam created this revision.
saiislam added reviewers: grokos, hfinkel, jdoerfert, JonChesterfield, ronlieb, 
ABataev, mdtoguchi, kbobrovs, sdmitriev, gregrodgers, kkwli0, dreachem, Tyker, 
jsjodin.
saiislam requested review of this revision.
Herald added a project: clang.
Herald added a subscriber: cfe-commits.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D93525

Files:
  clang/test/Driver/clang-offload-bundler.c
  clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

Index: clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp
===================================================================
--- clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp
+++ clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp
@@ -22,6 +22,8 @@
 #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"
@@ -30,6 +32,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Host.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
@@ -81,6 +84,7 @@
                        "  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));
@@ -124,6 +128,22 @@
   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:
@@ -145,7 +165,7 @@
   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.
@@ -317,7 +337,7 @@
     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,
@@ -480,7 +500,7 @@
 
   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();
@@ -674,7 +694,7 @@
     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;
 
@@ -757,6 +777,8 @@
     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")
@@ -916,6 +938,152 @@
   return Error::success();
 }
 
+static Archive::Kind getDefaultArchiveKindForHost() {
+  return Triple(sys::getDefaultTargetTriple()).isOSDarwin() ? Archive::K_DARWIN
+                                                            : Archive::K_GNU;
+}
+
+/// Returns arch specific extension for device object files in the generated
+/// device specific archive in archive unbundling mode
+static StringRef getDeviceFileExtension(StringRef Device) {
+  if (Device.contains("gfx"))
+    return ".bc";
+  if (Device.contains("sm_"))
+    return ".cubin";
+  else {
+    WithColor::warning() << "Could not determine extension for archive "
+                            "members, using \".o\"\n";
+    return ".o";
+  }
+}
+
+static StringRef removeExtension(StringRef FileName) {
+  return (FileName.contains(".")) ? FileName.rsplit('.').first : FileName;
+}
+
+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 ArchiveDevice,
+                               std::string OffloadTargetDevice) {
+  return !OffloadTargetDevice.empty() && OffloadTargetDevice == ArchiveDevice;
+}
+
+/// UnbundleArchive takes an archive file (".a") as input containing bundled
+/// object files, and an offload target (not host), and extracts the device code
+/// into a new archive file. The resulting archive file contains all device
+/// object files corresponding to given offload target device, with the same
+/// names, except the file extension have been modified depending on the target
+/// architecture (either ".cubin" or ".bc") The created archive file does not
+/// contain an index of the symbols.
+static Error UnbundleArchive() {
+  std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
+  std::vector<std::unique_ptr<std::string>> LibraryNames;
+  std::string OffloadTargetDevice = 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();
+
+    Optional<StringRef> OptionalCurKindTriple = *CurTripleOrErr;
+    // No device code in this child, skip
+    if (!OptionalCurKindTriple.hasValue())
+      continue;
+    StringRef CurKindTriple = *OptionalCurKindTriple;
+    assert(!CurKindTriple.empty());
+
+    while (!CurKindTriple.empty()) {
+      if (hasHostKind(CurKindTriple)) {
+        // Do nothing, we don't extract host code yet
+      } else if (checkDeviceOptions(getDevice(getTriple(CurKindTriple)),
+                                    OffloadTargetDevice)) {
+        std::string BundleData;
+        raw_string_ostream DataStream(BundleData);
+        if (Error Err = FileHandler.get()->ReadBundle(DataStream, *ChildBuffer))
+          return Err;
+
+        LibraryNames.push_back(std::unique_ptr<std::string>(new std::string(
+            getDeviceLibraryFileName(ChildName, OffloadTargetDevice))));
+        auto MemBuf = MemoryBuffer::getMemBufferCopy(
+            DataStream.str(), *(LibraryNames.back().get()));
+        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';
 }
@@ -949,13 +1117,25 @@
           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 (FilesType == "a") {
+      reportError(createStringError(errc::invalid_argument,
+                                    "Archive files are only supported "
+                                    "for unbundling"));
+      Error = true;
+    }
     if (OutputFileNames.size() != 1) {
       Error = true;
       reportError(createStringError(
@@ -1030,7 +1210,11 @@
   if (!llvm::sys::fs::exists(BundlerExecutable))
     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;
   }
Index: clang/test/Driver/clang-offload-bundler.c
===================================================================
--- clang/test/Driver/clang-offload-bundler.c
+++ clang/test/Driver/clang-offload-bundler.c
@@ -23,6 +23,7 @@
 //
 // RUN: echo 'Content of device file 1' > %t.tgt1
 // RUN: echo 'Content of device file 2' > %t.tgt2
+// RUN: echo 'Content of device file 3' > %t.tgt3
 
 //
 // Check help message.
@@ -90,6 +91,12 @@
 // CK-ERR9A: error: expecting exactly one host target but got 0
 // CK-ERR9B: error: expecting exactly one host target but got 2
 
+// RUN: not clang-offload-bundler -type=a -targets=hxst-powerpcxxle-ibm-linux-gnu,openxp-pxxerpc64le-ibm-linux-gnu,xpenmp-x86_xx-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.i 2>&1 | FileCheck %s --check-prefix CK-ERR10A
+// RUN: not clang-offload-bundler -type=a -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.i,%t.tgt1,%t.tgt2 -inputs=%t.bundle.i.notexist -unbundle 2>&1 | FileCheck %s -DFILE=%t.bundle.i.notexist --check-prefix CK-ERR10B
+
+// CK-ERR10A: error: Archive files are only supported for unbundling
+// CK-ERR10B: error: number of output files and targets should be 1 when unbundling an archive
+
 //
 // Check text bundle. This is a readable format, so we check for the format we expect to find.
 //
@@ -288,6 +295,41 @@
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
 
+// COM: Archive unbundling test #1:
+// COM: The input is an archive of bundled object file and the output should be
+// COM: an archive of device-specific files extracted from the bundled object
+// COM: files in the archive.
+
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.o,%t.tgt1 -outputs=%t.abundle1.o
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.o,%t.tgt2 -outputs=%t.abundle2.o
+// RUN: llvm-ar cr %t.lib.a %t.abundle1.o %t.abundle2.o
+// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.lib.a -outputs=%t.devicelib.a
+// RUN: llvm-ar t %t.devicelib.a | FileCheck %s
+// CHECK: abundle1.bc
+// CHECK: abundle2.bc
+
+// COM: Create an archive (*-archive.lib.a) containing two bundles of
+// COM: device 1 (gfx903) and one bundle of device 2 (gfx906)
+
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t.o,%t.tgt1 -outputs=%t.bundle-1-gfx903.o
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx906 -inputs=%t.o,%t.tgt1 -outputs=%t.bundle-1-gfx906.o
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t.o,%t.tgt2 -outputs=%t.bundle-2-gfx903.o
+// RUN: llvm-ar cr %t-archive.lib.a %t.bundle-1-gfx903.o %t.bundle-1-gfx906.o %t.bundle-2-gfx903.o
+
+// COM: Unbundle archive.lib.a to create a device sepcific archive for device 1
+// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t-archive.lib.a -outputs=%t-archive.lib-gfx903.a
+// RUN: llvm-ar t %t-archive.lib-gfx903.a | FileCheck %s -check-prefix=GFX903
+// GFX903: bundle-1-gfx903.bc
+// GFX903: bundle-2-gfx903.bc
+// GFX903-NOT: bundle-1-gfx906.bc
+
+// COM: Unbundle archive.lib.a to create a device sepcific archive for device 2
+// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx906 -inputs=%t-archive.lib.a -outputs=%t-archive.lib-gfx906.a
+// RUN: llvm-ar t %t-archive.lib-gfx906.a | FileCheck %s -check-prefix=GFX906
+// GFX906: bundle-1-gfx906.bc
+// GFX906-NOT: bundle-1-gfx903.bc
+// GFX906-NOT: bundle-2-gfx903.bc
+
 // Some code so that we can create a binary out of this file.
 int A = 0;
 void test_func(void) {
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to