Minimize use of ELF64LEObjectFile over ELFObjectFileBase in annotating_importer.

gematria/datasets/BUILD.bazel

@@ -72,6 +72,7 @@ cc_library(
         "@com_google_perf_data_converter//quipper:perf_data_cc_proto",
         "@com_google_perf_data_converter//quipper:perf_parser",
         "@com_google_perf_data_converter//quipper:perf_reader",
+        "@llvm-project//llvm:BinaryFormat",
         "@llvm-project//llvm:Object",
         "@llvm-project//llvm:Support",
     ],

gematria/datasets/annotating_importer.cc

@@ -95,67 +95,40 @@ absl::Status AnnotatingImporter::LoadBinary(std::string_view file_name) {
   return absl::OkStatus();
 }
 
-absl::StatusOr<llvm::object::ELF64LEObjectFile *>
+absl::StatusOr<llvm::object::ELFObjectFileBase *>
 AnnotatingImporter::GetELFFromBinary() {
   llvm::object::Binary *binary = owning_binary_.getBinary();
   if (!binary->isObject()) {
     return absl::InvalidArgumentError(
         absl::StrFormat("The given binary (%s) is not an object.",
                         std::string(binary->getFileName())));
   }
-  llvm::object::ObjectFile *object =
-      llvm::cast<llvm::object::ObjectFile>(binary);
-  if (!object) {
+  auto *object = llvm::cast<llvm::object::ObjectFile>(binary);
+  if (object == nullptr) {
     return absl::InvalidArgumentError(
         absl::StrFormat("Could not cast the binary (%s) to an ObjectFile.",
                         std::string(binary->getFileName())));
   }
 
   // Make sure the object is an ELF file.
-  if (!object->isELF() || !object->is64Bit() || !object->isLittleEndian()) {
+  if (!object->isELF()) {
     return absl::InvalidArgumentError(
-        absl::StrFormat("The given object (%s) is not in ELF64LE format.",
+        absl::StrFormat("The given object (%s) is not in ELF format.",
                         std::string(binary->getFileName())));
   }
-  auto *elf_object = llvm::dyn_cast<llvm::object::ELF64LEObjectFile>(object);
-  if (!elf_object) {
+  auto *elf_object = llvm::dyn_cast<llvm::object::ELFObjectFileBase>(object);
+  if (elf_object == nullptr) {
     return absl::InvalidArgumentError(absl::StrFormat(
-        "Could not cast the object (%s) to an ELF64LEObjectFile.",
+        "Could not cast the object (%s) to an ELFObjectFileBase.",
         std::string(binary->getFileName())));
   }
 
   return elf_object;
 }
 
-absl::StatusOr<llvm::object::Elf_Phdr_Impl<llvm::object::ELF64LE>>
-AnnotatingImporter::GetMainProgramHeader(
-    const llvm::object::ELF64LEObjectFile *elf_object) {
-  llvm::object::Elf_Phdr_Impl<llvm::object::ELF64LE> main_header;
-  bool found_main_header = false;
-  auto program_headers = elf_object->getELFFile().program_headers();
-  if (llvm::Error error = program_headers.takeError()) {
-    return LlvmErrorToStatus(std::move(error));
-  }
-  for (const auto &program_header : *program_headers) {
-    if (program_header.p_type == llvm::ELF::PT_LOAD &&
-        program_header.p_flags & llvm::ELF::PF_R &&
-        program_header.p_flags & llvm::ELF::PF_X) {
-      if (found_main_header) {
-        return absl::InvalidArgumentError(
-            "The given object has multiple executable segments. This is "
-            "currently not supported.");
-      }
-      main_header = program_header;
-      found_main_header = true;
-    }
-  }
-
-  return main_header;
-}
-
 absl::StatusOr<std::vector<DisassembledInstruction>>
 AnnotatingImporter::GetELFSlice(
-    const llvm::object::ELF64LEObjectFile *elf_object, uint64_t range_begin,
+    const llvm::object::ELFObjectFileBase *elf_object, uint64_t range_begin,
     uint64_t range_end, uint64_t file_offset) {
   llvm::StringRef binary_buf = elf_object->getData();
 
@@ -189,7 +162,24 @@ AnnotatingImporter::GetBlocksFromELF() {
   if (llvm::Error error = bb_addr_map.takeError()) {
     return LlvmErrorToStatus(std::move(error));
   }
-  const auto main_header = GetMainProgramHeader(*elf_object);
+
+  // TODO(vbshah): Consider making it possible to use other ELFTs rather than
+  // only ELF64LE since only the implementation of GetMainProgramHeader differs
+  // between different ELFTs.
+  if (!(*elf_object)->is64Bit() || !(*elf_object)->isLittleEndian()) {
+    return absl::InvalidArgumentError(
+        absl::StrFormat("The given object (%s) is not in ELF64LE format.",
+                        (*elf_object)->getFileName()));
+  }
+  auto *typed_elf_object =
+      llvm::dyn_cast<llvm::object::ELF64LEObjectFile>(*elf_object);
+  if (typed_elf_object == nullptr) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Could not cast the ELF object (%s) to an ELF64LEObjectFileBase.",
+        (*elf_object)->getFileName()));
+  }
+
+  const auto main_header = GetMainProgramHeader(typed_elf_object);
   if (!main_header.ok()) {
     return main_header.status();
   }
@@ -290,7 +280,24 @@ AnnotatingImporter::GetLBRBlocksWithLatency() {
   if (!elf_object.ok()) {
     return elf_object.status();
   }
-  const auto main_header = GetMainProgramHeader(*elf_object);
+
+  // TODO(vbshah): Consider making it possible to use other ELFTs rather than
+  // only ELF64LE since only the implementation of GetMainProgramHeader differs
+  // between different ELFTs.
+  if (!(*elf_object)->is64Bit() || !(*elf_object)->isLittleEndian()) {
+    return absl::InvalidArgumentError(
+        absl::StrFormat("The given object (%s) is not in ELF64LE format.",
+                        (*elf_object)->getFileName()));
+  }
+  auto *typed_elf_object =
+      llvm::dyn_cast<llvm::object::ELF64LEObjectFile>(*elf_object);
+  if (typed_elf_object == nullptr) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Could not cast the ELF object (%s) to an ELF64LEObjectFileBase.",
+        (*elf_object)->getFileName()));
+  }
+
+  const auto main_header = GetMainProgramHeader(typed_elf_object);
   if (!main_header.ok()) {
     return main_header.status();
   }

gematria/datasets/annotating_importer.h

@@ -29,10 +29,13 @@
 #include "gematria/datasets/bhive_importer.h"
 #include "gematria/llvm/canonicalizer.h"
 #include "gematria/llvm/disassembler.h"
+#include "gematria/llvm/llvm_to_absl.h"
 #include "gematria/proto/throughput.pb.h"
+#include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ELFTypes.h"
+#include "llvm/Support/Error.h"
 #include "quipper/perf_data.pb.h"
 #include "quipper/perf_parser.h"
 #include "quipper/perf_reader.h"
@@ -67,16 +70,17 @@ class AnnotatingImporter {
   // Returns a pointer inside the loaded binary casted down to an ELF object.
   // The pointer is owned by this instance of `AnnotatingImporter` and may only
   // be accessed while this is alive.
-  absl::StatusOr<llvm::object::ELF64LEObjectFile*> GetELFFromBinary();
+  absl::StatusOr<llvm::object::ELFObjectFileBase*> GetELFFromBinary();
 
-  // Returns the file offset of the passed ELF object.
-  absl::StatusOr<llvm::object::Elf_Phdr_Impl<llvm::object::ELF64LE>>
-  GetMainProgramHeader(const llvm::object::ELF64LEObjectFile* elf_object);
+  // Returns the program header corresponding to the main executable section.
+  template <class ELFT>
+  absl::StatusOr<llvm::object::Elf_Phdr_Impl<ELFT>> GetMainProgramHeader(
+      const llvm::object::ELFObjectFile<ELFT>* elf_object);
 
   // Disassembles and returns instructions between two addresses in an ELF
   // object.
   absl::StatusOr<std::vector<DisassembledInstruction>> GetELFSlice(
-      const llvm::object::ELF64LEObjectFile* elf_object, uint64_t range_begin,
+      const llvm::object::ELFObjectFileBase* elf_object, uint64_t range_begin,
       uint64_t range_end, uint64_t file_offset);
 
   // Extracts basic blocks from an ELF object, and returns them as tuple
@@ -112,6 +116,33 @@ class AnnotatingImporter {
   llvm::object::OwningBinary<llvm::object::Binary> owning_binary_;
 };
 
+template <class ELFT>
+absl::StatusOr<llvm::object::Elf_Phdr_Impl<ELFT>>
+AnnotatingImporter::GetMainProgramHeader(
+    const llvm::object::ELFObjectFile<ELFT>* elf_object) {
+  llvm::object::Elf_Phdr_Impl<ELFT> main_header;
+  bool found_main_header = false;
+  auto program_headers = elf_object->getELFFile().program_headers();
+  if (llvm::Error error = program_headers.takeError()) {
+    return LlvmErrorToStatus(std::move(error));
+  }
+  for (const auto& program_header : *program_headers) {
+    if (program_header.p_type == llvm::ELF::PT_LOAD &&
+        program_header.p_flags & llvm::ELF::PF_R &&
+        program_header.p_flags & llvm::ELF::PF_X) {
+      if (found_main_header) {
+        return absl::InvalidArgumentError(
+            "The given object has multiple executable segments. This is "
+            "currently not supported.");
+      }
+      main_header = program_header;
+      found_main_header = true;
+    }
+  }
+
+  return main_header;
+}
+
 }  // namespace gematria
 
 #endif  // THIRD_PARTY_GEMATRIA_GEMATRIA_DATASETS_ANNOTATING_IMPORTER_H_