Add llvm-extract-bundle-entry to extend llvm-objcopy

  This commit creates llvm-extract-bundle-entry as a wrapper to llvm-objcopy,
  to allow extracting HIP offload fatbin bundles given a URI argument.

Author: dsalinas_amdeng

llvm/include/llvm/Object/OffloadBundle.h

@@ -210,8 +210,8 @@ LLVM_ABI Error extractOffloadBundleFatBinary(
 
 /// Extract code object memory from the given \p Source object file at \p Offset
 /// and of \p Size, and copy into \p OutputFileName.
-LLVM_ABI Error extractCodeObject(const ObjectFile &Source, int64_t Offset,
-                                 int64_t Size, StringRef OutputFileName);
+LLVM_ABI Error extractCodeObject(const ObjectFile &Source, size_t Offset,
+                                 size_t Size, StringRef OutputFileName);
 
 /// Extract code object memory from the given \p Source object file at \p Offset
 /// and of \p Size, and copy into \p OutputFileName.

llvm/lib/Object/OffloadBundle.cpp

@@ -217,24 +217,37 @@ Error object::extractOffloadBundleFatBinary(
   return Error::success();
 }
 
-Error object::extractCodeObject(const ObjectFile &Source, int64_t Offset,
-                                int64_t Size, StringRef OutputFileName) {
+Error object::extractCodeObject(const ObjectFile &Source, size_t Offset,
+                                size_t Size, StringRef OutputFileName) {
   Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
       FileOutputBuffer::create(OutputFileName, Size);
 
-  if (!BufferOrErr)
-    return BufferOrErr.takeError();
+  if (auto EC = BufferOrErr.takeError())
+    return EC;
 
   Expected<MemoryBufferRef> InputBuffOrErr = Source.getMemoryBufferRef();
   if (Error Err = InputBuffOrErr.takeError())
-    return Err;
+    return createFileError(OutputFileName, std::move(Err));
+  ;
+
+  if (Size > InputBuffOrErr->getBufferSize())
+    return createStringError("size in URI is larger than source");
+
+  if (Offset > InputBuffOrErr->getBufferSize())
+    return createStringError(inconvertibleErrorCode(),
+                             "offset in URI is beyond the size of the source");
+
+  if (Offset + Size > InputBuffOrErr->getBufferSize())
+    return createStringError(
+        inconvertibleErrorCode(),
+        "offset + size in URI is beyond the size of the source");
 
   std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);
   std::copy(InputBuffOrErr->getBufferStart() + Offset,
             InputBuffOrErr->getBufferStart() + Offset + Size,
             Buf->getBufferStart());
   if (Error E = Buf->commit())
-    return E;
+    return createFileError(OutputFileName, std::move(E));
 
   return Error::success();
 }
@@ -259,6 +272,7 @@ Error object::extractOffloadBundleByURI(StringRef URIstr) {
   // create a URI object
   Expected<std::unique_ptr<OffloadBundleURI>> UriOrErr(
       OffloadBundleURI::createOffloadBundleURI(URIstr, FILE_URI));
+
   if (!UriOrErr)
     return UriOrErr.takeError();
 
@@ -275,7 +289,7 @@ Error object::extractOffloadBundleByURI(StringRef URIstr) {
   auto Obj = ObjOrErr->getBinary();
   if (Error Err =
           object::extractCodeObject(*Obj, Uri.Offset, Uri.Size, OutputFile))
-    return Err;
+    return createFileError(Uri.FileName, std::move(Err));
 
   return Error::success();
 }

llvm/test/CMakeLists.txt

@@ -103,6 +103,7 @@ set(LLVM_TEST_DEPENDS
   llvm-dwp
   llvm-exegesis
   llvm-extract
+  llvm-extract-bundle-entry
   llvm-gsymutil
   llvm-ir2vec
   llvm-isel-fuzzer

llvm/test/lit.cfg.py

@@ -219,6 +219,7 @@ def get_asan_rtlib():
     ToolSubst("%llvm-strip", FindTool("llvm-strip")),
     ToolSubst("%llvm-install-name-tool", FindTool("llvm-install-name-tool")),
     ToolSubst("%llvm-bitcode-strip", FindTool("llvm-bitcode-strip")),
+    ToolSubst("%llvm-extract-bundle-entry", FindTool("llvm-extract-bundle-entry")),
     ToolSubst("%split-file", FindTool("split-file")),
 ]
 
@@ -250,6 +251,7 @@ def get_asan_rtlib():
         "llvm-dlltool",
         "llvm-exegesis",
         "llvm-extract",
+        "llvm-extract-bundle-entry",
         "llvm-ir2vec",
         "llvm-isel-fuzzer",
         "llvm-ifs",

llvm/test/tools/llvm-objcopy/extract-bundle-entry.test

@@ -24,6 +24,10 @@ set(LLVM_TARGET_DEFINITIONS StripOpts.td)
 tablegen(LLVM StripOpts.inc -gen-opt-parser-defs)
 add_public_tablegen_target(StripOptsTableGen)
 
+set(LLVM_TARGET_DEFINITIONS ExtractBundleEntryOpts.td)
+tablegen(LLVM ExtractBundleEntryOpts.inc -gen-opt-parser-defs)
+add_public_tablegen_target(ExtractBundleEntryOptsTableGen)
+
 add_llvm_tool(llvm-objcopy
   ObjcopyOptions.cpp
   llvm-objcopy.cpp
@@ -37,10 +41,12 @@ add_llvm_tool(llvm-objcopy
 add_llvm_tool_symlink(llvm-install-name-tool llvm-objcopy)
 add_llvm_tool_symlink(llvm-bitcode-strip llvm-objcopy)
 add_llvm_tool_symlink(llvm-strip llvm-objcopy)
+add_llvm_tool_symlink(llvm-extract-bundle-entry llvm-objcopy)
 
 if(LLVM_INSTALL_BINUTILS_SYMLINKS)
   add_llvm_tool_symlink(objcopy llvm-objcopy)
   add_llvm_tool_symlink(strip llvm-objcopy)
+  add_llvm_tool_symlink(extract-bundle-entry llvm-objcopy)
 endif()
 
 if(LLVM_INSTALL_CCTOOLS_SYMLINKS)

llvm/tools/llvm-objcopy/CMakeLists.txt

@@ -0,0 +1,31 @@
+//===-- ExtractBundleEntryOpts.td - llvm-bitcode-strip options  ---------------*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file describes the command line options of llvm-extract-offload-entry.
+//
+//===----------------------------------------------------------------------===//
+
+include "llvm/Option/OptParser.td"
+
+def help : Flag<["--"], "Help">,
+           HelpText<"URIs can be read from STDIN, one per line.\n"
+                    "From the URIs specified, extracts code objects into files named:\n"
+                    "\t<executable_name>-[pid<number>]-offset<number>-size<number>.co\n\n"
+                    "URI syntax:\n" 
+                    "\tcode_object_uri ::== file_uri | memory_uri\n" 
+                    "\tfile_uri        ::== \"file://\" extract_file [ range_specifier ] \n" 
+                    "\tmemory_uri      ::== \"memory://\" process_id range_specifier\n" 
+                    "\trange_specifier ::== range_delimiter range_attribute [\"&\" range_attribute]\n" 
+                    "\trange_delimiter ::== \"#\" | \"?\"\n" 
+                    "\trange_attribute ::== [\"offset=\" number | \"size=\" number ]\n" 
+                    "\textract_file    ::== URI_ENCODED_OS_FILE_PATH\n" 
+                    "\tprocess_id      ::== DECIMAL_NUMBER\n" 
+                    "\tnumber          ::== HEX_NUMBER | DECIMAL_NUMBER | OCTAL_NUMBER\n" 
+                    "\nExample: file://dir1/dir2/hello_world#offset=133&size=14472 \n" 
+                    " memory://1234#offset=0x20000&size=3000\n">;
+def h : Flag<["-"], "h">, Alias<help>;

llvm/tools/llvm-objcopy/ExtractBundleEntryOpts.td

@@ -16,6 +16,7 @@
 #include "llvm/ObjCopy/ConfigManager.h"
 #include "llvm/ObjCopy/MachO/MachOConfig.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/OffloadBundle.h"
 #include "llvm/Option/Arg.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Support/CRC.h"
@@ -164,6 +165,41 @@ class StripOptTable : public opt::GenericOptTable {
   }
 };
 
+enum ExtractBundleEntryID {
+  EXTRACT_BUNDLE_ENTRY_INVALID = 0, // This is not an option ID.
+#define OPTION(...)                                                            \
+  LLVM_MAKE_OPT_ID_WITH_ID_PREFIX(EXTRACT_BUNDLE_ENTRY_, __VA_ARGS__),
+#include "ExtractBundleEntryOpts.inc"
+#undef OPTION
+};
+
+namespace extract_bundle_entry {
+#define OPTTABLE_STR_TABLE_CODE
+#include "ExtractBundleEntryOpts.inc"
+#undef OPTTABLE_STR_TABLE_CODE
+
+#define OPTTABLE_PREFIXES_TABLE_CODE
+#include "ExtractBundleEntryOpts.inc"
+#undef OPTTABLE_PREFIXES_TABLE_CODE
+
+static constexpr opt::OptTable::Info ExtractBundleEntryInfoTable[] = {
+#define OPTION(...)                                                            \
+  LLVM_CONSTRUCT_OPT_INFO_WITH_ID_PREFIX(EXTRACT_BUNDLE_ENTRY_, __VA_ARGS__),
+#include "ExtractBundleEntryOpts.inc"
+#undef OPTION
+};
+} // namespace extract_bundle_entry
+
+class ExtractBundleEntryOptTable : public opt::GenericOptTable {
+public:
+  ExtractBundleEntryOptTable()
+      : GenericOptTable(extract_bundle_entry::OptionStrTable,
+                        extract_bundle_entry::OptionPrefixesTable,
+                        extract_bundle_entry::ExtractBundleEntryInfoTable) {
+    setGroupedShortOptions(true);
+  }
+};
+
 } // namespace
 
 static SectionFlag parseSectionRenameFlag(StringRef SectionName) {
@@ -418,7 +454,13 @@ template <class T> static ErrorOr<T> getAsInteger(StringRef Val) {
 
 namespace {
 
-enum class ToolType { Objcopy, Strip, InstallNameTool, BitcodeStrip };
+enum class ToolType {
+  Objcopy,
+  Strip,
+  InstallNameTool,
+  BitcodeStrip,
+  ExtractBundleEntry
+};
 
 } // anonymous namespace
 
@@ -442,6 +484,10 @@ static void printHelp(const opt::OptTable &OptTable, raw_ostream &OS,
     ToolName = "llvm-bitcode-strip";
     HelpText = " [options] input";
     break;
+  case ToolType::ExtractBundleEntry:
+    ToolName = "llvm-extract-bundle-entry";
+    HelpText = " URI";
+    break;
   }
   OptTable.printHelp(OS, (ToolName + HelpText).str().c_str(),
                      (ToolName + " tool").str().c_str());
@@ -1670,3 +1716,41 @@ objcopy::parseStripOptions(ArrayRef<const char *> RawArgsArr,
 
   return std::move(DC);
 }
+
+Expected<DriverConfig> objcopy::parseExtractBundleEntryOptions(
+    ArrayRef<const char *> ArgsArr, function_ref<Error(Error)> ErrorCallback) {
+
+  DriverConfig DC;
+
+  ExtractBundleEntryOptTable T;
+  unsigned MissingArgumentIndex, MissingArgumentCount;
+  opt::InputArgList InputArgs =
+      T.ParseArgs(ArgsArr, MissingArgumentIndex, MissingArgumentCount);
+
+  if (InputArgs.size() == 0) {
+    printHelp(T, errs(), ToolType::ExtractBundleEntry);
+    exit(1);
+  }
+
+  if (InputArgs.hasArg(EXTRACT_BUNDLE_ENTRY_help)) {
+    printHelp(T, outs(), ToolType::ExtractBundleEntry);
+    exit(0);
+  }
+
+  for (auto *Arg : InputArgs.filtered(EXTRACT_BUNDLE_ENTRY_UNKNOWN))
+    return createStringError(errc::invalid_argument, "unknown argument '%s'",
+                             Arg->getAsString(InputArgs).c_str());
+
+  SmallVector<StringRef, 256> Positional;
+
+  for (auto *Arg : InputArgs.filtered(EXTRACT_BUNDLE_ENTRY_INPUT))
+    Positional.push_back(Arg->getValue());
+  assert(!Positional.empty());
+
+  // iterate over all input arguments
+  for (auto input : Positional)
+    if (Error Err = object::extractOffloadBundleByURI(input))
+      return std::move(Err);
+
+  return std::move(DC);
+}

llvm/tools/llvm-objcopy/ObjcopyOptions.cpp

@@ -51,6 +51,16 @@ parseBitcodeStripOptions(ArrayRef<const char *> ArgsArr,
 Expected<DriverConfig>
 parseStripOptions(ArrayRef<const char *> ArgsArr,
                   llvm::function_ref<Error(Error)> ErrorCallback);
+
+// ParseExtractBundleEntryOptions returns the config and sets the input
+// arguments. If a help flag is set then ParseExtractBundleEntryOptions will
+// print the help messege and exit. ErrorCallback is used to handle recoverable
+// errors. An Error returned by the callback aborts the parsing and is then
+// returned by this function.
+Expected<DriverConfig>
+parseExtractBundleEntryOptions(ArrayRef<const char *> ArgsArr,
+                               llvm::function_ref<Error(Error)> ErrorCallback);
+
 } // namespace objcopy
 } // namespace llvm
 

llvm/tools/llvm-objcopy/ObjcopyOptions.h

@@ -77,6 +77,8 @@ static Expected<DriverConfig> getDriverConfig(ArrayRef<const char *> Args) {
     return parseStripOptions(Args, reportWarning);
   else if (Is("install-name-tool") || Is("install_name_tool"))
     return parseInstallNameToolOptions(Args);
+  else if (Is("llvm-extract-bundle-entry"))
+    return parseExtractBundleEntryOptions(Args, reportWarning);
   else
     return parseObjcopyOptions(Args, reportWarning);
 }