jhuber6 updated this revision to Diff 453354.
jhuber6 added a comment.
Herald added a project: clang.
Herald added a subscriber: cfe-commits.

Moving to clang packager.


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
@@ -33,11 +33,20 @@
 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::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,26 +58,7 @@
   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;
-  }
-
-  auto reportError = [argv](Error E) {
-    logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
-    return EXIT_FAILURE;
-  };
-
+static Error bundleImages() {
   SmallVector<char, 1024> BinaryData;
   raw_svector_ostream OS(BinaryData);
   for (StringRef Image : DeviceImages) {
@@ -77,18 +67,17 @@
 
     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);
+      auto [Key, Value] = Arg.split("=");
+      if (Args.count(Key))
+        Args[Key] = Saver.save(Args[Key] + "," + Value);
       else
-        Args[KeyAndValue.first] = KeyAndValue.second;
+        Args[Key] = Value;
     }
 
     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;
@@ -98,7 +87,7 @@
         llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
             llvm::MemoryBuffer::getFileOrSTDIN(KeyAndValue.getValue());
         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)
@@ -115,18 +104,146 @@
     }
     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);
+
+    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;
+    }
+
+    for (const auto &Binary : *BinariesOrErr) {
+      // 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() + "." +
+                           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);
+      std::copy(Binary->getImage().bytes_begin(),
+                Binary->getImage().bytes_end(), 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;
 }
Index: clang/test/Driver/offload-packager.c
===================================================================
--- /dev/null
+++ clang/test/Driver/offload-packager.c
@@ -0,0 +1,16 @@
+// REQUIRES: x86-registered-target
+// REQUIRES: nvptx-registered-target
+// REQUIRES: amdgpu-registered-target
+
+// 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
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to