Fix offload extract (#10257)

- More comments for maintainability.
- Extract relevant sections up front to speed up the embedded binary
search.
- Provide better diagnostics (especially providing info on the result
files).
- Fix code to work on Windows (COFF files).

---------

Co-authored-by: Ejjeh <[email protected]>

Author: Dirk MG Seynhaeve

clang/test/Driver/clang-offload-extract.c

@@ -34,8 +34,8 @@
 // Extract target images (deprecated use model)
 //
 // RUN: clang-offload-extract --output=%t.deprecated %t.fat.bin | FileCheck %s --check-prefix CHECK-DEPRECATE
-// CHECK-DEPRECATE: Saving target image to
-// CHECK-DEPRECATE: Saving target image to
+// CHECK-DEPRECATE: Section {{.*}}: Image 1'-> File
+// CHECK-DEPRECATE: Section {{.*}}: Image 2'-> File
 
 //
 // Check that extracted contents match the original images.
@@ -47,8 +47,8 @@
 // Extract target images (new use model)
 //
 // RUN: clang-offload-extract --stem=%t.extracted %t.fat.bin | FileCheck %s --check-prefix CHECK-EXTRACT
-// CHECK-EXTRACT: Saving target image to
-// CHECK-EXTRACT: Saving target image to
+// CHECK-EXTRACT: Section {{.*}}: Image 1'-> File
+// CHECK-EXTRACT: Section {{.*}}: Image 2'-> File
 
 //
 // Check that extracted contents match the original images.

clang/tools/clang-offload-extract/ClangOffloadExtract.cpp

@@ -28,18 +28,24 @@
 #include "llvm/Support/raw_ostream.h"
 
 #define IMAGE_INFO_SECTION_NAME ".tgtimg"
+#define IMAGE_SECTION_NAME_PREFIX "__CLANG_OFFLOAD_BUNDLE__"
+// Windows truncates the names of sections to 8 bytes
+#define IMAGE_SECTION_NAME_PREFIX_COFF "__CLANG_"
+
+#define DEBUG_TYPE "clang-offload-extract"
 
 using namespace llvm;
 using namespace llvm::object;
 
+// Command-line parsing
 // Create a category to label utility-specific options; This will allow
 // us to distinguish those specific options from generic options and
 // irrelevant options
 static cl::OptionCategory
     ClangOffloadExtractCategory("Utility-specific options");
 
-// Create options for the input and output files, each with an
-// appropriate default when not specified
+// Create options for the input (fat binary) and output (target images)
+// files, each with an appropriate default when not specified
 static cl::opt<std::string> Input(cl::Positional, cl::init("a.out"),
                                   cl::desc("<input file>"),
                                   cl::cat(ClangOffloadExtractCategory));
@@ -61,10 +67,8 @@ static cl::alias FileNameStemAlias(
     "output", cl::desc("Deprecated option, replaced by option '--stem'"),
     cl::aliasopt(FileNameStem), cl::cat(ClangOffloadExtractCategory));
 
-// Path to the current binary
-static std::string ToolPath;
-
 // Report error (and handle any deferred errors)
+// This is the main diagnostic used by the utility
 static void reportError(Error E, Twine message = "\n") {
   std::string S;
   raw_string_ostream OSS(S);
@@ -80,7 +84,6 @@ static void reportError(Error E, Twine message = "\n") {
 
 int main(int argc, const char **argv) {
   sys::PrintStackTraceOnErrorSignal(argv[0]);
-  ToolPath = argv[0];
 
   // Hide non-generic options that are not in this utility's explicit
   // category;
@@ -138,105 +141,144 @@ linked fat binary, and store them in separate files.
   }
 
   // We are dealing with an appropriate fat binary;
-  // Locate the section IMAGE_INFO_SECTION_NAME (which contains the
-  // metadata on the embedded binaries)
-  unsigned FileNum = 0;
-
+  // * Create an array all the sections that have
+  //   IMAGE_SECTION_NAME in the section name:
+  auto OffloadSections = SmallVector<SectionRef>();
+  // * Locate the section that starts with IMAGE_INFO_SECTION_NAME_PREFIX
+  //   and extract the index for all the embedded binaries:
+  auto OffloadIndex = SmallVector<SectionRef>();
   for (const auto &Section : Binary->sections()) {
-    Expected<StringRef> NameOrErr = Section.getName();
-    if (auto E = NameOrErr.takeError()) {
+    Expected<StringRef> InfoSecNameOrErr = Section.getName();
+    if (auto E = InfoSecNameOrErr.takeError()) {
       reportError(std::move(E), "Input File: '" + Input + "'\n");
     }
-    if (*NameOrErr != IMAGE_INFO_SECTION_NAME)
-      continue;
+    LLVM_DEBUG(dbgs() << "Section: " << *InfoSecNameOrErr << "\n");
 
-    // This is the section we are looking for;
-    // Extract the meta information:
-    // The IMAGE_INFO_SECTION_NAME section contains packed <address,
-    // size> pairs describing target images that are stored in the fat
-    // binary.
-    Expected<StringRef> DataOrErr = Section.getContents();
-    if (auto E = DataOrErr.takeError()) {
-      reportError(std::move(E), "Input File: '" + Input + "'\n");
+    // We have a valid section name
+    std::string SectionNameToCompare = isa<COFFObjectFile>(Binary)
+                                           ? IMAGE_SECTION_NAME_PREFIX_COFF
+                                           : IMAGE_SECTION_NAME_PREFIX;
+    if (InfoSecNameOrErr->find(SectionNameToCompare) != std::string::npos) {
+      // This section contains embedded binaries
+      OffloadSections.push_back(Section);
+    } else if (*InfoSecNameOrErr == IMAGE_INFO_SECTION_NAME) {
+      // This section is the index
+      OffloadIndex.push_back(Section);
     }
-    // Data type to store the metadata for an individual target image
-    struct ImgInfoType {
-      uintptr_t Addr;
-      uintptr_t Size;
-    };
-
-    // Store the metadata for all target images in an array of target
-    // image information descriptors
-    auto ImgInfo =
-        ArrayRef(reinterpret_cast<const ImgInfoType *>(DataOrErr->data()),
-                 DataOrErr->size() / sizeof(ImgInfoType));
-
-    //  Loop over the image information descriptors to extract each
-    // target image.
-    for (const auto &Img : ImgInfo) {
-      // Ignore zero padding that can be inserted by the linker.
-      if (!Img.Addr)
-        continue;
-
-      // Find section which contains this image.
-      // TODO: can use more efficient algorithm than linear search. For
-      // example sections and images could be sorted by address then one pass
-      // performed through both at the same time.
-      // std::find_if
-      // * searches for a true predicate in [first,last] =~ [first,end)
-      // * returns end if no predicate is true
-      // It is probably faster to track  success through a bool (ImgFound)
-      bool ImgFound = false;
-      auto ImgSec =
-          find_if(Binary->sections(), [&Img, &ImgFound](SectionRef Sec) {
-            bool pred = (                         //
-                Sec.isData()                      //
-                && (Img.Addr == Sec.getAddress()) //
-                && (Img.Size == Sec.getSize())    //
-            );
-            ImgFound = ImgFound || pred;
-            return pred;
-          });
-      if (!ImgFound) {
-
-        reportError(
-            createStringError(
-                inconvertibleErrorCode(),
-                "cannot find section containing <0x%lx, 0x%lx> target image",
-                Img.Addr, Img.Size),
-            "Input File: '" + Input + "'\n");
-      }
-
-      Expected<StringRef> SecDataOrErr = ImgSec->getContents();
-      if (auto E = SecDataOrErr.takeError()) {
-        reportError(std::move(E), "Input File: '" + Input + "'\n");
-      }
-
-      // Output file name is composed from the name prefix provided by the
-      // user and the image number which is appended to the prefix
-      std::string FileName = FileNameStem + "." + std::to_string(FileNum++);
-
-      // Tell user that we are saving an image.
-      outs() << "Saving target image to '" << FileName << "'\n";
-
-      // Write image data to the output
-      std::error_code EC;
-      raw_fd_ostream OS(FileName, EC);
-      if (EC) {
-        reportError(createFileError(FileName, EC),
-                    "Specify a different Output File ('--stem' option)\n");
-      }
-
-      OS << SecDataOrErr->substr(Img.Addr - ImgSec->getAddress(), Img.Size);
-      if (OS.has_error()) {
-        reportError(createFileError(FileName, OS.error()),
-                    "Try a different Output File ('--stem' option)");
-      }
-    } // &Img: ImgInfo
-
-    // Fat binary is not expected to have more than one .tgtimg section.
-    break;
   }
 
+  // Check if there are any sections with embedded binaries
+  if (OffloadSections.size() == 0) {
+    reportError(
+        createStringError(inconvertibleErrorCode(),
+                          "Could not locate sections with offload binaries"),
+        "Fat Binary: '" + Input + "'\n");
+  }
+  // Check if we found the index section
+  if (OffloadIndex.size() == 0) {
+    reportError(
+        createStringError(inconvertibleErrorCode(),
+                          "Could not locate index for embedded binaries"),
+        "Fat Binary: '" + Input + "'\n");
+  }
+  // Check if we have a valid index section
+  Expected<StringRef> DataOrErr = OffloadIndex[0].getContents();
+  if (auto E = DataOrErr.takeError()) {
+    reportError(std::move(E), "Input File: '" + Input + "'\n");
+  }
+
+  // The input binary contains embedded offload binaries
+  // The index section contains packed <address, size> pairs describing
+  // target images that are stored in the fat binary.
+  // Data type to store the index for an individual target image
+  struct ImgInfoType {
+    uintptr_t Addr;
+    uintptr_t Size;
+  };
+
+  // Store the metadata for all target images in an array of target
+  // image information descriptors
+  // This can be done by reinterpreting the content of the section
+  auto ImgInfo =
+      ArrayRef(reinterpret_cast<const ImgInfoType *>(DataOrErr->data()),
+               DataOrErr->size() / sizeof(ImgInfoType));
+
+  //  Loop over the image information descriptors to extract each
+  // target image the object file data
+  unsigned FileNum = 0;
+  unsigned ImgCnt = 1;
+
+  for (const auto &Img : ImgInfo) {
+    // Ignore zero padding that can be inserted by the linker.
+    if (!Img.Addr)
+      continue;
+
+    // Find section which contains this image.
+    // /!\ There might be multiple images in a section
+    // std::find_if
+    // * searches for a true predicate in [first,last] =~ [first,end)
+    // * returns end if no predicate is true
+    // It is probably faster to track  success through a bool (ImgFound)
+    bool ImgFound = false;
+    auto ImgSec = find_if(OffloadSections, [&Img, &ImgFound](auto Sec) {
+      bool pred = (                                                        //
+          Sec.isData()                                                     //
+          && (Img.Addr >= Sec.getAddress())                                //
+          && ((Img.Addr + Img.Size) <= (Sec.getAddress() + Sec.getSize())) //
+      );
+      ImgFound = ImgFound || pred;
+      return pred;
+    });
+    if (!ImgFound) {
+
+      reportError(
+          createStringError(inconvertibleErrorCode(),
+                            "Target image (Address:0x%lx, Size:0x%lx>) is "
+                            "not contained in any section",
+                            Img.Addr, Img.Size),
+          "Fat Binary: '" + Input + "'\n");
+    }
+
+    Expected<StringRef> ImgSecNameOrErr = ImgSec->getName();
+    if (auto E = ImgSecNameOrErr.takeError()) {
+      reportError(std::move(E),
+                  "Can not determine section name in Fat Binary '" + Input +
+                      "'\n");
+    }
+    Expected<StringRef> SecDataOrErr = ImgSec->getContents();
+    if (auto E = SecDataOrErr.takeError()) {
+      reportError(std::move(E), "Can not extract contents for section '" +
+                                    *ImgSecNameOrErr + "' in Fat Binary '" +
+                                    Input + "'\n");
+    }
+
+    // Output file name is composed from the name prefix provided by the
+    // user and the image number which is appended to the prefix
+    std::string FileName = FileNameStem + "." + std::to_string(FileNum++);
+    std::string OffloadName = ImgSecNameOrErr->data();
+    std::string OffloadPrefix = isa<COFFObjectFile>(Binary)
+                                    ? IMAGE_SECTION_NAME_PREFIX_COFF
+                                    : IMAGE_SECTION_NAME_PREFIX;
+    OffloadName.erase(0, OffloadPrefix.length());
+
+    // Tell user that we are saving an image.
+    outs() << "Section '" + OffloadName + "': Image " << ImgCnt++
+           << "'-> File '" + FileName + "'\n";
+
+    // Write image data to the output
+    std::error_code EC;
+    raw_fd_ostream OS(FileName, EC);
+    if (EC) {
+      reportError(createFileError(FileName, EC),
+                  "Specify a different Output File ('--stem' option)\n");
+    }
+
+    OS << SecDataOrErr->substr(Img.Addr - ImgSec->getAddress(), Img.Size);
+    if (OS.has_error()) {
+      reportError(createFileError(FileName, OS.error()),
+                  "Try a different Output File ('--stem' option)");
+    }
+  } // &Img: ImgInfo
+
   return 0;
 }