[ORC] Support scanning "fallback" slices for interfaces.

llvm/include/llvm/ExecutionEngine/Orc/MachO.h

@@ -95,21 +95,36 @@ class ForceLoadMachOArchiveMembers {
   bool ObjCOnly;
 };
 
+using GetFallbackArchsFn =
+    unique_function<SmallVector<std::pair<uint32_t, uint32_t>>(
+        uint32_t CPUType, uint32_t CPUSubType)>;
+
+/// Match the exact CPU type/subtype only.
+LLVM_ABI SmallVector<std::pair<uint32_t, uint32_t>>
+noFallbackArchs(uint32_t CPUType, uint32_t CPUSubType);
+
+/// Match standard dynamic loader fallback rules.
+LLVM_ABI SmallVector<std::pair<uint32_t, uint32_t>>
+standardMachOFallbackArchs(uint32_t CPUType, uint32_t CPUSubType);
+
 /// Returns a SymbolNameSet containing the exported symbols defined in the
 /// given dylib.
-LLVM_ABI Expected<SymbolNameSet>
-getDylibInterfaceFromDylib(ExecutionSession &ES, Twine Path);
+LLVM_ABI Expected<SymbolNameSet> getDylibInterfaceFromDylib(
+    ExecutionSession &ES, Twine Path,
+    GetFallbackArchsFn GetFallbackArchs = standardMachOFallbackArchs);
 
 /// Returns a SymbolNameSet containing the exported symbols defined in the
 /// relevant slice of the TapiUniversal file.
-LLVM_ABI Expected<SymbolNameSet>
-getDylibInterfaceFromTapiFile(ExecutionSession &ES, Twine Path);
+LLVM_ABI Expected<SymbolNameSet> getDylibInterfaceFromTapiFile(
+    ExecutionSession &ES, Twine Path,
+    GetFallbackArchsFn GetFallbackArchs = standardMachOFallbackArchs);
 
 /// Returns a SymbolNameSet containing the exported symbols defined in the
 /// relevant slice of the given file, which may be either a dylib or a tapi
 /// file.
-LLVM_ABI Expected<SymbolNameSet> getDylibInterface(ExecutionSession &ES,
-                                                   Twine Path);
+LLVM_ABI Expected<SymbolNameSet> getDylibInterface(
+    ExecutionSession &ES, Twine Path,
+    GetFallbackArchsFn GetFallbackArchs = standardMachOFallbackArchs);
 
 } // namespace orc
 } // namespace llvm

llvm/lib/ExecutionEngine/Orc/MachO.cpp

@@ -282,15 +282,48 @@ Expected<bool> ForceLoadMachOArchiveMembers::operator()(
   return true;
 }
 
-Expected<SymbolNameSet> getDylibInterfaceFromDylib(ExecutionSession &ES,
-                                                   Twine Path) {
-  auto CPUType = MachO::getCPUType(ES.getTargetTriple());
-  if (!CPUType)
-    return CPUType.takeError();
+LLVM_ABI SmallVector<std::pair<uint32_t, uint32_t>>
+noFallbackArchs(uint32_t CPUType, uint32_t CPUSubType) {
+  SmallVector<std::pair<uint32_t, uint32_t>> Result;
+  Result.push_back({CPUType, CPUSubType});
+  return Result;
+}
+
+SmallVector<std::pair<uint32_t, uint32_t>>
+standardMachOFallbackArchs(uint32_t CPUType, uint32_t CPUSubType) {
+  SmallVector<std::pair<uint32_t, uint32_t>> Archs;
+
+  // Match given CPU type/subtype first.
+  Archs.push_back({CPUType, CPUSubType});
+
+  switch (CPUType) {
+  case MachO::CPU_TYPE_ARM64:
+    // Handle arm64 variants.
+    switch (CPUSubType) {
+    case MachO::CPU_SUBTYPE_ARM64_ALL:
+      Archs.push_back({CPUType, MachO::CPU_SUBTYPE_ARM64E});
+      break;
+    default:
+      break;
+    }
+    break;
+  default:
+    break;
+  }
+
+  return Archs;
+}
+
+Expected<SymbolNameSet>
+getDylibInterfaceFromDylib(ExecutionSession &ES, Twine Path,
+                           GetFallbackArchsFn GetFallbackArchs) {
+  auto InitCPUType = MachO::getCPUType(ES.getTargetTriple());
+  if (!InitCPUType)
+    return InitCPUType.takeError();
 
-  auto CPUSubType = MachO::getCPUSubType(ES.getTargetTriple());
-  if (!CPUSubType)
-    return CPUSubType.takeError();
+  auto InitCPUSubType = MachO::getCPUSubType(ES.getTargetTriple());
+  if (!InitCPUSubType)
+    return InitCPUSubType.takeError();
 
   auto Buf = MemoryBuffer::getFile(Path);
   if (!Buf)
@@ -301,25 +334,38 @@ Expected<SymbolNameSet> getDylibInterfaceFromDylib(ExecutionSession &ES,
     return BinFile.takeError();
 
   std::unique_ptr<object::MachOObjectFile> MachOFile;
-  if (isa<object::MachOObjectFile>(**BinFile))
+  if (isa<object::MachOObjectFile>(**BinFile)) {
     MachOFile.reset(dyn_cast<object::MachOObjectFile>(BinFile->release()));
-  else if (auto *MachOUni =
-               dyn_cast<object::MachOUniversalBinary>(BinFile->get())) {
-    for (auto &O : MachOUni->objects()) {
-      if (O.getCPUType() == *CPUType &&
-          (O.getCPUSubType() & ~MachO::CPU_SUBTYPE_MASK) == *CPUSubType) {
-        if (auto Obj = O.getAsObjectFile())
-          MachOFile = std::move(*Obj);
-        else
-          return Obj.takeError();
-        break;
+
+    // TODO: Check that dylib arch is compatible.
+  } else if (auto *MachOUni =
+                 dyn_cast<object::MachOUniversalBinary>(BinFile->get())) {
+    SmallVector<std::pair<uint32_t, uint32_t>> ArchsToTry;
+    if (GetFallbackArchs)
+      ArchsToTry = GetFallbackArchs(*InitCPUType, *InitCPUSubType);
+    else
+      ArchsToTry.push_back({*InitCPUType, *InitCPUSubType});
+
+    for (auto &[CPUType, CPUSubType] : ArchsToTry) {
+      for (auto &O : MachOUni->objects()) {
+        if (O.getCPUType() == CPUType &&
+            (O.getCPUSubType() & ~MachO::CPU_SUBTYPE_MASK) == CPUSubType) {
+          if (auto Obj = O.getAsObjectFile())
+            MachOFile = std::move(*Obj);
+          else
+            return Obj.takeError();
+          break;
+        }
       }
+      if (MachOFile) // If found, break out.
+        break;
     }
     if (!MachOFile)
-      return make_error<StringError>("MachO universal binary at " + Path +
-                                         " does not contain a slice for " +
-                                         ES.getTargetTriple().str(),
-                                     inconvertibleErrorCode());
+      return make_error<StringError>(
+          "MachO universal binary at " + Path +
+              " does not contain a compatible slice for " +
+              ES.getTargetTriple().str(),
+          inconvertibleErrorCode());
   } else
     return make_error<StringError>("File at " + Path + " is not a MachO",
                                    inconvertibleErrorCode());
@@ -339,8 +385,9 @@ Expected<SymbolNameSet> getDylibInterfaceFromDylib(ExecutionSession &ES,
   return std::move(Symbols);
 }
 
-Expected<SymbolNameSet> getDylibInterfaceFromTapiFile(ExecutionSession &ES,
-                                                      Twine Path) {
+Expected<SymbolNameSet>
+getDylibInterfaceFromTapiFile(ExecutionSession &ES, Twine Path,
+                              GetFallbackArchsFn GetFallbackArchs) {
   SymbolNameSet Symbols;
 
   auto TapiFileBuffer = MemoryBuffer::getFile(Path);
@@ -352,37 +399,54 @@ Expected<SymbolNameSet> getDylibInterfaceFromTapiFile(ExecutionSession &ES,
   if (!Tapi)
     return Tapi.takeError();
 
-  auto CPUType = MachO::getCPUType(ES.getTargetTriple());
-  if (!CPUType)
-    return CPUType.takeError();
+  auto InitCPUType = MachO::getCPUType(ES.getTargetTriple());
+  if (!InitCPUType)
+    return InitCPUType.takeError();
 
-  auto CPUSubType = MachO::getCPUSubType(ES.getTargetTriple());
-  if (!CPUSubType)
-    return CPUSubType.takeError();
+  auto InitCPUSubType = MachO::getCPUSubType(ES.getTargetTriple());
+  if (!InitCPUSubType)
+    return InitCPUSubType.takeError();
+
+  SmallVector<std::pair<uint32_t, uint32_t>> ArchsToTry;
+  if (GetFallbackArchs)
+    ArchsToTry = GetFallbackArchs(*InitCPUType, *InitCPUSubType);
+  else
+    ArchsToTry.push_back({*InitCPUType, *InitCPUSubType});
 
   auto &IF = (*Tapi)->getInterfaceFile();
-  auto Interface =
-      IF.extract(MachO::getArchitectureFromCpuType(*CPUType, *CPUSubType));
-  if (!Interface)
-    return Interface.takeError();
 
-  for (auto *Sym : (*Interface)->exports())
-    Symbols.insert(ES.intern(Sym->getName()));
+  auto ArchSet = IF.getArchitectures();
+  for (auto [CPUType, CPUSubType] : ArchsToTry) {
+    auto A = MachO::getArchitectureFromCpuType(CPUType, CPUSubType);
+    if (ArchSet.has(A)) {
+      if (auto Interface = IF.extract(A)) {
+        for (auto *Sym : (*Interface)->exports())
+          Symbols.insert(ES.intern(Sym->getName()));
+        return Symbols;
+      } else
+        return Interface.takeError();
+    }
+  }
 
-  return Symbols;
+  return make_error<StringError>(
+      "MachO interface file at " + Path +
+          " does not contain a compatible slice for " +
+          ES.getTargetTriple().str(),
+      inconvertibleErrorCode());
 }
 
-Expected<SymbolNameSet> getDylibInterface(ExecutionSession &ES, Twine Path) {
+Expected<SymbolNameSet> getDylibInterface(ExecutionSession &ES, Twine Path,
+                                          GetFallbackArchsFn GetFallbackArchs) {
   file_magic Magic;
   if (auto EC = identify_magic(Path, Magic))
     return createFileError(Path, EC);
 
   switch (Magic) {
   case file_magic::macho_universal_binary:
   case file_magic::macho_dynamically_linked_shared_lib:
-    return getDylibInterfaceFromDylib(ES, Path);
+    return getDylibInterfaceFromDylib(ES, Path, std::move(GetFallbackArchs));
   case file_magic::tapi_file:
-    return getDylibInterfaceFromTapiFile(ES, Path);
+    return getDylibInterfaceFromTapiFile(ES, Path, std::move(GetFallbackArchs));
   default:
     return make_error<StringError>("Cannot get interface for " + Path +
                                        " unrecognized file type",

llvm/test/ExecutionEngine/JITLink/AArch64/Inputs/MachO_Foo_arm64e.tbd

@@ -0,0 +1,23 @@
+--- !tapi-tbd
+tbd-version: 4
+targets:  [ arm64e-macos ]
+uuids:
+  - target: arm64e-macos
+    value: 00000000-0000-0000-0000-000000000000
+flags: [ installapi ]
+install-name: Foo.framework/Foo
+current-version: 1.2.3
+compatibility-version: 1.2
+swift-abi-version: 5
+parent-umbrella:
+  - targets: [ arm64e-macos ]
+    umbrella: System
+exports:
+  - targets: [ arm64e-macos ]
+    symbols: [ _foo ]
+    objc-classes: []
+    objc-eh-types: []
+    objc-ivars: []
+    weak-symbols: []
+    thread-local-symbols: []
+...

llvm/test/ExecutionEngine/JITLink/AArch64/MachO_weak_link.test

@@ -1,8 +1,14 @@
 # RUN: rm -rf %t && mkdir -p %t
 # RUN: llvm-mc -triple=arm64-apple-darwin19 -filetype=obj -o %t/main.o \
 # RUN:         %S/Inputs/MachO_main_ret_foo.s
-# RUN: llvm-jitlink -noexec %t/main.o -weak_library %S/Inputs/MachO_Foo.tbd
-
+# RUN: llvm-jitlink -noexec %t/main.o -weak_library \
+# RUN:     %S/Inputs/MachO_Foo_arm64.tbd
+# RUN: llvm-jitlink -noexec %t/main.o -weak_library \
+# RUN:     %S/Inputs/MachO_Foo_arm64e.tbd
+#
 # Check that we can load main.o, which unconditionally uses symbol foo, by
 # using -weak_library on a TBD file to emulate forced weak linking against
 # a library that supplies foo, but is missing at runtime.
+#
+# Check that weak linking works for arm64 JIT'd programs even if the TBD
+# file contains only an arm64e interface.