[ORC] loadRelocatableObject: universal binary support, clearer errors (#104406)

lhames · web-flow · commit 3e1d4ec671c5 · 2024-08-16T09:52:36.000+10:00
ORC supports loading relocatable object files into a JIT'd process. The
raw "add object file" API (ObjectLayer::add) accepts plain relocatable
object files as llvm::MemoryBuffers only and does not check that the
object file's format or architecture are compatible with the process
that it will be linked in to. This API is flexible, but places the
burden of error checking and universal binary support on clients.

This commit introduces a new utility, loadRelocatableObject, that takes
a path to load and a target triple and then:
1. If the path does not exist, returns a FileError containing the
invalid path.
2. If the path points to a MachO universal binary, identifies and
returns MemoryBuffer covering the slice that matches the given triple
(checking that the slice really does contains a valid MachO relocatable
object with a compatible arch).
3. If the path points to a regular relocatable object file, verifies
that the format and architecture are compatible with the triple.

Clients can use loadRelocatableObject in the common case of loading
object files from disk to simplify their code.

Note: Error checking for ELF and COFF is left as a FIXME.

rdar://133653290
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -319,9 +319,6 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
                                    Error &Err);
   Error buildObjectFilesMap();
 
-  static Expected<std::pair<size_t, size_t>>
-  getSliceRangeForArch(object::MachOUniversalBinary &UB, const Triple &TT);
-
   ObjectLayer &L;
   GetObjectFileInterface GetObjFileInterface;
   std::set<std::string> ImportedDynamicLibraries;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LoadRelocatableObject.h b/llvm/include/llvm/ExecutionEngine/Orc/LoadRelocatableObject.h
@@ -0,0 +1,36 @@
+//===---- LoadRelocatableObject.h - Load relocatable objects ----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A wrapper for `MemoryBuffer::getFile` / `MemoryBuffer::getFileSlice` that:
+//
+//   1. Adds file paths to errors by default.
+//   2. Checks architecture compatibility up-front.
+//   3. Handles MachO universal binaries, returning the MemoryBuffer for the
+//      requested slice only.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LOADRELOCATABLEOBJECT_H
+#define LLVM_EXECUTIONENGINE_ORC_LOADRELOCATABLEOBJECT_H
+
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/TargetParser/Triple.h"
+
+namespace llvm {
+namespace orc {
+
+// Load an object file compatible with the given triple (if given) from the
+// given path. May return a file slice if the path contains a universal binary.
+Expected<std::unique_ptr<MemoryBuffer>> loadRelocatableObject(StringRef Path,
+                                                              const Triple &TT);
+
+} // End namespace orc
+} // End namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LOADRELOCATABLEOBJECT_H
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MachO.h b/llvm/include/llvm/ExecutionEngine/Orc/MachO.h
@@ -0,0 +1,63 @@
+//===------------- MachO.h - MachO format utilities -------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains utilities for load MachO relocatable object files.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_MACHO_H
+#define LLVM_EXECUTIONENGINE_ORC_MACHO_H
+
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/TargetParser/Triple.h"
+
+namespace llvm {
+
+namespace object {
+
+class MachOUniversalBinary;
+
+} // namespace object
+
+namespace orc {
+
+/// Check that the given buffer contains a MachO object file compatible with the
+/// given triple.
+/// ObjIsSlice should be set to true if Obj is a slice of a universal binary
+/// (that fact will then be reported in the error messages).
+Expected<std::unique_ptr<MemoryBuffer>>
+checkMachORelocatableObject(std::unique_ptr<MemoryBuffer> Obj, const Triple &TT,
+                            bool ObjIsSlice);
+
+/// Load a relocatable object compatible with TT from Path.
+/// If Path is a universal binary, this function will return a buffer for the
+/// slice compatible with Triple (if one is present).
+Expected<std::unique_ptr<MemoryBuffer>>
+loadMachORelocatableObject(StringRef Path, const Triple &TT);
+
+/// Load a compatible relocatable object (if available) from a MachO universal
+/// binary.
+Expected<std::unique_ptr<MemoryBuffer>>
+loadMachORelocatableObjectFromUniversalBinary(
+    StringRef UBPath, std::unique_ptr<MemoryBuffer> UBBuf, const Triple &TT);
+
+/// Utility for identifying the file-slice compatible with TT in a universal
+/// binary.
+Expected<std::pair<size_t, size_t>>
+getMachOSliceRangeForTriple(object::MachOUniversalBinary &UB, const Triple &TT);
+
+/// Utility for identifying the file-slice compatible with TT in a universal
+/// binary.
+Expected<std::pair<size_t, size_t>>
+getMachOSliceRangeForTriple(MemoryBufferRef UBBuf, const Triple &TT);
+
+} // namespace orc
+} // namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_MACHO_H
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -29,8 +29,10 @@ add_llvm_component_library(LLVMOrcJIT
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
+  LoadRelocatableObject.cpp
   LookupAndRecordAddrs.cpp
   LLJIT.cpp
+  MachO.cpp
   MachOPlatform.cpp
   MapperJITLinkMemoryManager.cpp
   MemoryMapper.cpp
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -9,6 +9,7 @@
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/MachO.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
@@ -295,7 +296,7 @@ StaticLibraryDefinitionGenerator::Load(
 
     const auto &TT = L.getExecutionSession().getTargetTriple();
 
-    auto SliceRange = getSliceRangeForArch(*UB, TT);
+    auto SliceRange = getMachOSliceRangeForTriple(*UB, TT);
     if (!SliceRange)
       return SliceRange.takeError();
 
@@ -358,7 +359,7 @@ StaticLibraryDefinitionGenerator::Create(
 
     const auto &TT = L.getExecutionSession().getTargetTriple();
 
-    auto SliceRange = getSliceRangeForArch(*UB, TT);
+    auto SliceRange = getMachOSliceRangeForTriple(*UB, TT);
     if (!SliceRange)
       return SliceRange.takeError();
 
@@ -460,27 +461,6 @@ Error StaticLibraryDefinitionGenerator::buildObjectFilesMap() {
   return Error::success();
 }
 
-Expected<std::pair<size_t, size_t>>
-StaticLibraryDefinitionGenerator::getSliceRangeForArch(
-    object::MachOUniversalBinary &UB, const Triple &TT) {
-
-  for (const auto &Obj : UB.objects()) {
-    auto ObjTT = Obj.getTriple();
-    if (ObjTT.getArch() == TT.getArch() &&
-        ObjTT.getSubArch() == TT.getSubArch() &&
-        (TT.getVendor() == Triple::UnknownVendor ||
-         ObjTT.getVendor() == TT.getVendor())) {
-      // We found a match. Return the range for the slice.
-      return std::make_pair(Obj.getOffset(), Obj.getSize());
-    }
-  }
-
-  return make_error<StringError>(Twine("Universal binary ") + UB.getFileName() +
-                                     " does not contain a slice for " +
-                                     TT.str(),
-                                 inconvertibleErrorCode());
-}
-
 StaticLibraryDefinitionGenerator::StaticLibraryDefinitionGenerator(
     ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
     std::unique_ptr<object::Archive> Archive,
diff --git a/llvm/lib/ExecutionEngine/Orc/LoadRelocatableObject.cpp b/llvm/lib/ExecutionEngine/Orc/LoadRelocatableObject.cpp
@@ -0,0 +1,69 @@
+//===----- LoadRelocatableObject.cpp -- Load relocatable object files -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LoadRelocatableObject.h"
+#include "llvm/BinaryFormat/Magic.h"
+#include "llvm/ExecutionEngine/Orc/MachO.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm {
+namespace orc {
+
+static Expected<std::unique_ptr<MemoryBuffer>>
+checkCOFFRelocatableObject(std::unique_ptr<MemoryBuffer> Obj,
+                           const Triple &TT) {
+  // TODO: Actually check the architecture of the file.
+  return std::move(Obj);
+}
+
+static Expected<std::unique_ptr<MemoryBuffer>>
+checkELFRelocatableObject(std::unique_ptr<MemoryBuffer> Obj, const Triple &TT) {
+  // TODO: Actually check the architecture of the file.
+  return std::move(Obj);
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+loadRelocatableObject(StringRef Path, const Triple &TT) {
+  auto Buf = MemoryBuffer::getFile(Path);
+  if (!Buf)
+    return createFileError(Path, Buf.getError());
+
+  std::optional<Triple::ObjectFormatType> RequireFormat;
+  if (TT.getObjectFormat() != Triple::UnknownObjectFormat)
+    RequireFormat = TT.getObjectFormat();
+
+  switch (identify_magic((*Buf)->getBuffer())) {
+  case file_magic::coff_object:
+    if (!RequireFormat || *RequireFormat == Triple::COFF)
+      return checkCOFFRelocatableObject(std::move(*Buf), TT);
+    break;
+  case file_magic::elf_relocatable:
+    if (!RequireFormat || *RequireFormat == Triple::ELF)
+      return checkELFRelocatableObject(std::move(*Buf), TT);
+    break;
+  case file_magic::macho_object:
+    if (!RequireFormat || *RequireFormat == Triple::MachO)
+      return checkMachORelocatableObject(std::move(*Buf), TT, false);
+    break;
+  case file_magic::macho_universal_binary:
+    if (!RequireFormat || *RequireFormat == Triple::MachO)
+      return loadMachORelocatableObjectFromUniversalBinary(Path,
+                                                           std::move(*Buf), TT);
+    break;
+  default:
+    break;
+  }
+  return make_error<StringError>(
+      Path + " does not contain a relocatable object file compatible with " +
+          TT.str(),
+      inconvertibleErrorCode());
+}
+
+} // End namespace orc.
+} // End namespace llvm.
diff --git a/llvm/lib/ExecutionEngine/Orc/MachO.cpp b/llvm/lib/ExecutionEngine/Orc/MachO.cpp
diff --git a/llvm/test/ExecutionEngine/JITLink/Generic/MachO_universal_binaries.test b/llvm/test/ExecutionEngine/JITLink/Generic/MachO_universal_binaries.test
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
