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

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

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -319,9 +319,6 @@ private:
                                    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;

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

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

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

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,

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.

llvm/lib/ExecutionEngine/Orc/MachO.cpp

@@ -0,0 +1,164 @@
+//===----------------- MachO.cpp - MachO format utilities -----------------===//
+//
+// 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/MachO.h"
+
+#include "llvm/BinaryFormat/MachO.h"
+#include "llvm/Object/MachOUniversal.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm {
+namespace orc {
+
+static std::string objDesc(MemoryBuffer &Obj, const Triple &TT,
+                           bool ObjIsSlice) {
+  std::string Desc;
+  if (ObjIsSlice)
+    Desc += (TT.getArchName() + " slice of universal binary").str();
+  Desc += Obj.getBufferIdentifier();
+  return Desc;
+}
+
+template <typename HeaderType>
+static Expected<std::unique_ptr<MemoryBuffer>>
+checkMachORelocatableObject(std::unique_ptr<MemoryBuffer> Obj,
+                            bool SwapEndianness, const Triple &TT,
+                            bool ObjIsSlice) {
+  StringRef Data = Obj->getBuffer();
+
+  HeaderType Hdr;
+  memcpy(&Hdr, Data.data(), sizeof(HeaderType));
+
+  if (SwapEndianness)
+    swapStruct(Hdr);
+
+  if (Hdr.filetype != MachO::MH_OBJECT)
+    return make_error<StringError>(objDesc(*Obj, TT, ObjIsSlice) +
+                                       " is not a MachO relocatable object",
+                                   inconvertibleErrorCode());
+
+  auto ObjArch = object::MachOObjectFile::getArch(Hdr.cputype, Hdr.cpusubtype);
+  if (ObjArch != TT.getArch())
+    return make_error<StringError>(
+        objDesc(*Obj, TT, ObjIsSlice) + Triple::getArchTypeName(ObjArch) +
+            ", cannot be loaded into " + TT.str() + " process",
+        inconvertibleErrorCode());
+
+  return std::move(Obj);
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+checkMachORelocatableObject(std::unique_ptr<MemoryBuffer> Obj, const Triple &TT,
+                            bool ObjIsSlice) {
+  StringRef Data = Obj->getBuffer();
+
+  if (Data.size() < 4)
+    return make_error<StringError>(
+        objDesc(*Obj, TT, ObjIsSlice) +
+            " is not a valid MachO relocatable object file (truncated header)",
+        inconvertibleErrorCode());
+
+  uint32_t Magic;
+  memcpy(&Magic, Data.data(), sizeof(uint32_t));
+
+  switch (Magic) {
+  case MachO::MH_MAGIC:
+  case MachO::MH_CIGAM:
+    return checkMachORelocatableObject<MachO::mach_header>(
+        std::move(Obj), Magic == MachO::MH_CIGAM, TT, ObjIsSlice);
+  case MachO::MH_MAGIC_64:
+  case MachO::MH_CIGAM_64:
+    return checkMachORelocatableObject<MachO::mach_header_64>(
+        std::move(Obj), Magic == MachO::MH_CIGAM_64, TT, ObjIsSlice);
+  default:
+    return make_error<StringError>(
+        objDesc(*Obj, TT, ObjIsSlice) +
+            " is not a valid MachO relocatable object (bad magic value)",
+        inconvertibleErrorCode());
+  }
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+loadMachORelocatableObject(StringRef Path, const Triple &TT) {
+  assert((TT.getObjectFormat() == Triple::UnknownObjectFormat ||
+          TT.getObjectFormat() == Triple::MachO) &&
+         "TT must specify MachO or Unknown object format");
+
+  auto Buf = MemoryBuffer::getFile(Path);
+  if (!Buf)
+    return createFileError(Path, Buf.getError());
+
+  switch (identify_magic((*Buf)->getBuffer())) {
+  case file_magic::macho_object:
+    return checkMachORelocatableObject(std::move(*Buf), TT, false);
+  case file_magic::macho_universal_binary:
+    return loadMachORelocatableObjectFromUniversalBinary(Path, std::move(*Buf),
+                                                         TT);
+  default:
+    return make_error<StringError>(
+        Path + " does not contain a relocatable object file compatible with " +
+            TT.str(),
+        inconvertibleErrorCode());
+  }
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+loadMachORelocatableObjectFromUniversalBinary(
+    StringRef UBPath, std::unique_ptr<MemoryBuffer> UBBuf, const Triple &TT) {
+
+  auto UniversalBin =
+      object::MachOUniversalBinary::create(UBBuf->getMemBufferRef());
+  if (!UniversalBin)
+    return UniversalBin.takeError();
+
+  auto SliceRange = getMachOSliceRangeForTriple(**UniversalBin, TT);
+  if (!SliceRange)
+    return SliceRange.takeError();
+
+  auto ObjBuf = errorOrToExpected(MemoryBuffer::getFileSlice(
+      UBPath, SliceRange->second, SliceRange->first, false));
+  if (!ObjBuf)
+    return createFileError(UBPath, ObjBuf.takeError());
+
+  return checkMachORelocatableObject(std::move(*ObjBuf), TT, true);
+}
+
+Expected<std::pair<size_t, size_t>>
+getMachOSliceRangeForTriple(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());
+}
+
+Expected<std::pair<size_t, size_t>>
+getMachOSliceRangeForTriple(MemoryBufferRef UBBuf, const Triple &TT) {
+
+  auto UB = object::MachOUniversalBinary::create(UBBuf);
+  if (!UB)
+    return UB.takeError();
+
+  return getMachOSliceRangeForTriple(**UB, TT);
+}
+
+} // End namespace orc.
+} // End namespace llvm.

llvm/test/ExecutionEngine/JITLink/Generic/MachO_universal_binaries.test

@@ -0,0 +1,9 @@
+# RUN: llc -filetype=obj -mtriple x86_64-apple-macosx -o %t.x86-64.o \
+# RUN:   %S/Inputs/main-ret-0.ll
+# RUN: llc -filetype=obj -mtriple arm64-apple-macosx -o %t.arm64.o \
+# RUN:   %S/Inputs/main-ret-0.ll
+# RUN: llvm-lipo -create -output %t.universal.o -arch x86_64 %t.x86-64.o \
+# RUN:   -arch arm64 %t.arm64.o
+# RUN: llvm-jitlink -noexec -triple arm64-apple-macosx %t.universal.o
+#
+# Check that universal binaries don't raise "not a valid object file" errors.

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -28,6 +28,8 @@
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/LoadRelocatableObject.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
@@ -260,6 +262,10 @@ static cl::opt<bool> UseSharedMemory(
     cl::desc("Use shared memory to transfer generated code and data"),
     cl::init(false), cl::cat(JITLinkCategory));
 
+static cl::opt<std::string>
+    OverrideTriple("triple", cl::desc("Override target triple detection"),
+                   cl::init(""), cl::cat(JITLinkCategory));
+
 static ExitOnError ExitOnErr;
 
 static LLVM_ATTRIBUTE_USED void linkComponents() {
@@ -1101,7 +1107,8 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
   for (auto &HarnessFile : TestHarnesses) {
     HarnessFiles.insert(HarnessFile);
 
-    auto ObjBuffer = ExitOnErr(getFile(HarnessFile));
+    auto ObjBuffer =
+        ExitOnErr(loadRelocatableObject(HarnessFile, ES.getTargetTriple()));
 
     auto ObjInterface =
         ExitOnErr(getObjectFileInterface(ES, ObjBuffer->getMemBufferRef()));
@@ -1418,6 +1425,14 @@ Session::findSymbolInfo(StringRef SymbolName, Twine ErrorMsgStem) {
 static std::pair<Triple, SubtargetFeatures> getFirstFileTripleAndFeatures() {
   static std::pair<Triple, SubtargetFeatures> FirstTTAndFeatures = []() {
     assert(!InputFiles.empty() && "InputFiles can not be empty");
+
+    if (!OverrideTriple.empty()) {
+      LLVM_DEBUG({
+        dbgs() << "Triple from -triple override: " << OverrideTriple << "\n";
+      });
+      return std::make_pair(Triple(OverrideTriple), SubtargetFeatures());
+    }
+
     for (auto InputFile : InputFiles) {
       auto ObjBuffer = ExitOnErr(getFile(InputFile));
       file_magic Magic = identify_magic(ObjBuffer->getBuffer());
@@ -1436,13 +1451,25 @@ static std::pair<Triple, SubtargetFeatures> getFirstFileTripleAndFeatures() {
         SubtargetFeatures Features;
         if (auto ObjFeatures = Obj->getFeatures())
           Features = std::move(*ObjFeatures);
+
+        LLVM_DEBUG({
+          dbgs() << "Triple from " << InputFile << ": " << TT.str() << "\n";
+        });
         return std::make_pair(TT, Features);
       }
       default:
         break;
       }
     }
-    return std::make_pair(Triple(), SubtargetFeatures());
+
+    // If no plain object file inputs exist to pin down the triple then detect
+    // the host triple and default to that.
+    auto JTMB = ExitOnErr(JITTargetMachineBuilder::detectHost());
+    LLVM_DEBUG({
+      dbgs() << "Triple from host-detection: " << JTMB.getTargetTriple().str()
+             << "\n";
+    });
+    return std::make_pair(JTMB.getTargetTriple(), JTMB.getFeatures());
   }();
 
   return FirstTTAndFeatures;
@@ -1711,9 +1738,9 @@ static Error addSectCreates(Session &S,
                                          ", filename component cannot be empty",
                                      inconvertibleErrorCode());
 
-    auto Content = MemoryBuffer::getFile(FileName);
+    auto Content = getFile(FileName);
     if (!Content)
-      return createFileError(FileName, errorCodeToError(Content.getError()));
+      return Content.takeError();
 
     SectCreateMaterializationUnit::ExtraSymbolsMap ExtraSymbols;
     while (!ExtraSymbolsString.empty()) {
@@ -1745,7 +1772,7 @@ static Error addTestHarnesses(Session &S) {
   LLVM_DEBUG(dbgs() << "Adding test harness objects...\n");
   for (auto HarnessFile : TestHarnesses) {
     LLVM_DEBUG(dbgs() << "  " << HarnessFile << "\n");
-    auto ObjBuffer = getFile(HarnessFile);
+    auto ObjBuffer = loadRelocatableObject(HarnessFile, S.ES.getTargetTriple());
     if (!ObjBuffer)
       return ObjBuffer.takeError();
     if (auto Err = S.ObjLayer.add(*S.MainJD, std::move(*ObjBuffer)))
@@ -1770,7 +1797,7 @@ static Error addObjects(Session &S,
     auto &JD = *std::prev(IdxToJD.lower_bound(InputFileArgIdx))->second;
     LLVM_DEBUG(dbgs() << "  " << InputFileArgIdx << ": \"" << InputFile
                       << "\" to " << JD.getName() << "\n";);
-    auto ObjBuffer = getFile(InputFile);
+    auto ObjBuffer = loadRelocatableObject(InputFile, S.ES.getTargetTriple());
     if (!ObjBuffer)
       return ObjBuffer.takeError();