[lld] Implement safe icf for MachO

[alx32]

Summary

This change implements --icf=safe for MachO based on addrsig section that is implemented in D123751.

[thakis]

Very cool!

This needs a test case.

[int3]

Re testing: you can look at https://github.com/llvm/llvm-project/blob/main/lld/test/MachO/icf.s for inspiration. I would probably create a new test file that has some address-significant functions, and check that we don't dedup them with --icf=safe while continuing to do so with --icf=all.

[int3]

lld/test/MachO/icf-safe.s

10–14

if you add --macho to the objdump command (i.e. llvm-objdump --macho -d) these will be rendered as bl _func0<N> instead, then it's easy to check that those values are the same / different

[int3]

lld/MachO/Driver.cpp

1546–1548

nit: can omit if braces

lld/MachO/ICF.cpp

370–374

can omit braces

379–380

hmm if I understand correctly, the failure case we are trying to avoid here is that a DSO may export two symbols a and b, and whatever loads that DSO may end up comparing the addresses of a and b for equality? If that's the case I think any exported symbol in a DSO should be marked as address-significant. But... that seems like we would end up marking a lot of symbols :/

referencedDynamically is only used to mark symbols that shouldn't be stripped (and currently only applies to __mh_execute_header, as per its block comment); but stripped symbols can still have their addresses taken if they are in the export trie.

What does LLD-ELF do here?

382

I think the _or_null isn't necessary, this should always be nonnull

389

398–399

codebase convention is to avoid auto for simple types

424

lld/MachO/ICF.h

21

symtab and config are globals in LLD; no need to pass them explicitly

lld/MachO/InputSection.h

74

would be nice to have a comment here

lld/test/MachO/icf-options.s

12

can just reuse DIAG-EMPTY here, no need for DIAG-SAFE

lld/test/MachO/icf-safe.s

17–19

nit: align

35–40

we usually try to create test cases that are as minimal as possible. can this just be a ret void?

105–107

ditto, I suppose we need noinline but can we omit the others?

[int3]

lld/MachO/ICF.cpp

379–380

we should cover this with a test too

[int3]

lgtm!

lld/MachO/ICF.h

12

could forward-declare InputFile and move this include to the .cpp

20

oops, I belatedly realized that inputFiles is a global as well, doesn't need to be passed either :D

(we could do away with the SetVector.h include in that case)

lld/test/MachO/icf-safe.s

10–12

FileCheck operates via substring matching by default, so there's no need for the regex

nit: also align the lines with the _callAllFunctions line above

24

should remove the other attributes we deleted below (or just remove the entire comment altogether)

[smeenai]

lld/MachO/InputSection.h

74

Nit: "symbol" is a bit confusing, since the flag is on a section. Perhaps something like this instead?

[smeenai]

lld/test/MachO/icf-safe.s

25

Nit: you can inline the attributes instead of declaring them out of line (same for everything below)

lld/MachO/Driver.cpp

auto icfLevel = StringSwitch<ICFLevel>(icfLevelStr)
                       .Cases("none", "", ICFLevel::none)
                       .Case("safe", ICFLevel::safe)
                       .Case("all", ICFLevel::all)
                       .Default(ICFLevel::unknown);
   if (icfLevel == ICFLevel::unknown) {
     warn(Twine("unknown --icf=OPTION `") + icfLevelStr +
          "', defaulting to `none'");
     icfLevel = ICFLevel::none;
-  } else if (icfLevel == ICFLevel::safe) {
-    warn(Twine("`--icf=safe' is not yet implemented, reverting to `none'"));
-    icfLevel = ICFLevel::none;
   }
   return icfLevel;
 }
 
 static void warnIfDeprecatedOption(const Option &opt) {
   if (!opt.getGroup().isValid())
     return;
   if (opt.getGroup().getID() == OPT_grp_deprecated) {
if (config->callGraphProfileSort)
       priorityBuilder.extractCallGraphProfile();
 
     if (config->deadStrip)
       markLive();
 
     // ICF assumes that all literals have been folded already, so we must run
     // foldIdenticalLiterals before foldIdenticalSections.
     foldIdenticalLiterals();
-    if (config->icfLevel != ICFLevel::none)
+    if (config->icfLevel != ICFLevel::none) {
+      if (config->icfLevel == ICFLevel::safe)
+        markAddrSigSymbols();
       foldIdenticalSections();
+    }
 
     // Write to an output file.
     if (target->wordSize == 8)
       writeResult<LP64>();
     else
       writeResult<ILP32>();
 
     depTracker->write(getLLDVersion(), inputFiles, config->outputFile);

lld/MachO/ICF.h

 //===- ICF.h ----------------------------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLD_MACHO_ICF_H
 #define LLD_MACHO_ICF_H
 
+#include "InputFiles.h"
 #include "lld/Common/LLVM.h"
 #include <vector>
 
 namespace lld {
 namespace macho {
+class Symbol;
 
+void markAddrSigSymbols();
+void markSymAsAddrSig(Symbol *s);
 void foldIdenticalSections();
 
 } // namespace macho
 } // namespace lld
 
 #endif

lld/MachO/ICF.cpp

 //===- ICF.cpp ------------------------------------------------------------===//
 //
 // 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 "ICF.h"
 #include "ConcatOutputSection.h"
+#include "Config.h"
 #include "InputSection.h"
+#include "SymbolTable.h"
 #include "Symbols.h"
 #include "UnwindInfoSection.h"
 
 #include "lld/Common/CommonLinkerContext.h"
+#include "llvm/Support/LEB128.h"
 #include "llvm/Support/Parallel.h"
 #include "llvm/Support/TimeProfiler.h"
 #include "llvm/Support/xxhash.h"
 
 #include <atomic>
 
 using namespace llvm;
 using namespace lld;
while (begin < end) {
     // If we created a group, we need to iterate the main loop again.
     if (mid != end)
       icfRepeat = true;
 
     begin = mid;
   }
 }
 
+void macho::markSymAsAddrSig(Symbol *s) {
+  if (auto *d = dyn_cast_or_null<Defined>(s))
+    if (d->isec)
+      d->isec->keepUnique = true;
+}
+
+void macho::markAddrSigSymbols() {
+  for (InputFile *file : inputFiles) {
+    ObjFile *obj = dyn_cast<ObjFile>(file);
+    if (!obj)
+      continue;
+
+    Section *addrSigSection = obj->addrSigSection;
+    if (!addrSigSection)
+      continue;
+    assert(addrSigSection->subsections.size() == 1);
+
+    Subsection *subSection = &addrSigSection->subsections[0];
+    ArrayRef<unsigned char> &contents = subSection->isec->data;
+
+    const uint8_t *pData = contents.begin();
+    while (pData != contents.end()) {
+      unsigned size;
+      const char *err;
+      uint32_t symIndex = decodeULEB128(pData, &size, contents.end(), &err);
+      if (err)
+        fatal(toString(file) + ": could not decode addrsig section: " + err);
+      markSymAsAddrSig(obj->symbols[symIndex]);
+      pData += size;
+    }
+  }
+}
+
 void macho::foldIdenticalSections() {
   TimeTraceScope timeScope("Fold Identical Code Sections");
   // The ICF equivalence-class segregation algorithm relies on pre-computed
   // hashes of InputSection::data for the ConcatOutputSection::inputs and all
   // sections referenced by their relocs. We could recursively traverse the
   // relocs to find every referenced InputSection, but that precludes easy
   // parallelization. Therefore, we hash every InputSection here where we have
   // them all accessible as simple vectors.
 
   // If an InputSection is ineligible for ICF, we give it a unique ID to force
   // it into an unfoldable singleton equivalence class.  Begin the unique-ID
   // space at inputSections.size(), so that it will never intersect with
   // equivalence-class IDs which begin at 0. Since hashes & unique IDs never
   // coexist with equivalence-class IDs, this is not necessary, but might help
   // someone keep the numbers straight in case we ever need to debug the
   // ICF::segregate()
   std::vector<ConcatInputSection *> hashable;
   uint64_t icfUniqueID = inputSections.size();
   for (ConcatInputSection *isec : inputSections) {
     // FIXME: consider non-code __text sections as hashable?
     bool isHashable = (isCodeSection(isec) || isCfStringSection(isec) ||
                        isClassRefsSection(isec)) &&
-                      !isec->shouldOmitFromOutput() &&
+                      !isec->keepUnique && !isec->shouldOmitFromOutput() &&
                       sectionType(isec->getFlags()) == MachO::S_REGULAR;
     if (isHashable) {
       hashable.push_back(isec);
       for (Defined *d : isec->symbols)
         if (d->unwindEntry)
           hashable.push_back(d->unwindEntry);
 
       // __cfstring has embedded addends that foil ICF's hashing / equality

lld/MachO/InputFiles.h

public:
   template <class LP> void parse();
 
   static bool classof(const InputFile *f) { return f->kind() == ObjKind; }
 
   llvm::DWARFUnit *compileUnit = nullptr;
   const uint32_t modTime;
   std::vector<ConcatInputSection *> debugSections;
   std::vector<CallGraphEntry> callGraph;
+  Section *addrSigSection = nullptr;
 
 private:
   template <class LP> void parseLazy();
   template <class SectionHeader> void parseSections(ArrayRef<SectionHeader>);
   template <class LP>
   void parseSymbols(ArrayRef<typename LP::section> sectionHeaders,
                     ArrayRef<typename LP::nlist> nList, const char *strtab,
                     bool subsectionsViaSymbols);

lld/MachO/InputFiles.cpp

if (sectionType(sec.flags) == S_CSTRING_LITERALS ||
         checkError(parseCallGraph(data, callGraph));
       // ld64 does not appear to emit contents from sections within the __LLVM
       // segment. Symbols within those sections point to bitcode metadata
       // instead of actual symbols. Global symbols within those sections could
       // have the same name without causing duplicate symbol errors. To avoid
       // spurious duplicate symbol errors, we do not parse these sections.
       // TODO: Evaluate whether the bitcode metadata is needed.
     } else {
+      if (name == section_names::addrSig)
+        addrSigSection = sections.back();
+
       auto *isec = make<ConcatInputSection>(section, data, align);
       if (isDebugSection(isec->getFlags()) &&
           isec->getSegName() == segment_names::dwarf) {
         // Instead of emitting DWARF sections, we emit STABS symbols to the
         // object files that contain them. We filter them out early to avoid
         // parsing their relocations unnecessarily.
         debugSections.push_back(isec);
       } else {

lld/MachO/InputSection.h

InputSection(const InputSection &rhs)
       : sectionKind(rhs.sectionKind), align(rhs.align), data(rhs.data),
         section(rhs.section) {}
 
   Kind sectionKind;
 
 public:
   // is address assigned?
   bool isFinal = false;
+  // keep the address of the symbol(s) in this section unique in the final binary ?
+  bool keepUnique = false;
   uint32_t align = 1;
 
   OutputSection *parent = nullptr;
   ArrayRef<uint8_t> data;
   std::vector<Reloc> relocs;
   // The symbols that belong to this InputSection, sorted by value. With
   // .subsections_via_symbols, there is typically only one element here.
   llvm::TinyPtrVector<Defined *> symbols;

 constexpr const char symbolTable[] = "__symbol_table";
 constexpr const char textCoalNt[] = "__textcoal_nt";
 constexpr const char text[] = "__text";
 constexpr const char threadPtrs[] = "__thread_ptrs";
 constexpr const char threadVars[] = "__thread_vars";
 constexpr const char unwindInfo[] = "__unwind_info";
 constexpr const char weakBinding[] = "__weak_binding";
 constexpr const char zeroFill[] = "__zerofill";
+constexpr const char addrSig[] = "__llvm_addrsig";
 
 } // namespace section_names
 
 } // namespace macho
 
 std::string toString(const macho::InputSection *);
 
 } // namespace lld
 
 #endif

lld/test/MachO/icf-options.s

 # REQUIRES: x86
 # RUN: rm -rf %t; mkdir %t
 
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t/main.o
 # RUN: %lld -lSystem --icf=all -o %t/all %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
 # RUN: %lld -lSystem --icf=none -o %t/none %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
 # RUN: %lld -lSystem -no_deduplicate -o %t/no_dedup %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
-# RUN: not %lld -lSystem --icf=safe -o %t/safe %t/main.o 2>&1 \
-# RUN:     | FileCheck %s --check-prefix=DIAG-SAFE
+# RUN: %lld -lSystem --icf=safe -o %t/safe %t/main.o 2>&1 \
+# RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
 # RUN: not %lld -lSystem --icf=junk -o %t/junk %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-JUNK
 # RUN: %lld -lSystem --icf=all -no_deduplicate -o %t/none2 %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
 # RUN: %lld -lSystem -no_deduplicate --icf=all -o %t/all2 %t/main.o 2>&1 \
 # RUN:     | FileCheck %s --check-prefix=DIAG-EMPTY --allow-empty
 
 # DIAG-EMPTY-NOT: {{.}}
-# DIAG-SAFE: `--icf=safe' is not yet implemented, reverting to `none'
 # DIAG-JUNK: unknown --icf=OPTION `junk', defaulting to `none'
 
 # RUN: llvm-objdump -d --syms %t/all | FileCheck %s --check-prefix=FOLD
 # RUN: llvm-objdump -d --syms %t/all2 | FileCheck %s --check-prefix=FOLD
 # RUN: llvm-objdump -d --syms %t/none | FileCheck %s --check-prefix=NOOP
 # RUN: llvm-objdump -d --syms %t/none2 | FileCheck %s --check-prefix=NOOP
 # RUN: llvm-objdump -d --syms %t/no_dedup | FileCheck %s --check-prefix=NOOP
 

lld/test/MachO/icf-safe.s

+; RUN: rm -rf %t; mkdir %t
+; RUN: llc -filetype=obj %s -O3 -o %t/icf-obj.o -enable-machine-outliner=never -mtriple arm64-apple-ios -addrsig
+; RUN: %lld -arch arm64 -lSystem --icf=safe -dylib -o %t/icf-safe.dylib %t/icf-obj.o
+; RUN: %lld -arch arm64 -lSystem --icf=all  -dylib -o %t/icf-all.dylib  %t/icf-obj.o
+; RUN: llvm-objdump %t/icf-safe.dylib -d --macho | FileCheck %s --check-prefix=ICFSAFE
+; RUN: llvm-objdump %t/icf-all.dylib  -d --macho | FileCheck %s --check-prefix=ICFALL
+
+; ICFSAFE-LABEL:  _callAllFunctions
+; ICFSAFE:        bl _func02
+; ICFSAFE-NEXT:   bl _func02
+; ICFSAFE-NEXT:   bl _func03_takeaddr
+
+; ICFALL-LABEL:   _callAllFunctions
+; ICFALL:         bl _func03_takeaddr
+; ICFALL-NEXT:    bl _func03_takeaddr
+; ICFALL-NEXT:    bl _func03_takeaddr
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-ios7.0.0"
+
+@result = global i32 0, align 4
+
+define void @func01() local_unnamed_addr noinline {
+entry:
+  %0 = load volatile i32, ptr @result, align 4
+  %add = add nsw i32 %0, 1
+  store volatile i32 %add, ptr @result, align 4
+  ret void
+}
+
+define void @func02() local_unnamed_addr noinline {
+entry:
+  %0 = load volatile i32, ptr @result, align 4
+  %add = add nsw i32 %0, 1
+  store volatile i32 %add, ptr @result, align 4
+  ret void
+}
+
+define void @func03_takeaddr() noinline {
+entry:
+  %0 = load volatile i32, ptr @result, align 4
+  %add = add nsw i32 %0, 1
+  store volatile i32 %add, ptr @result, align 4
+  ret void
+}
+
+define void @callAllFunctions() local_unnamed_addr {
+entry:
+  tail call void @func01()
+  tail call void @func02()
+  tail call void @func03_takeaddr()
+  %0 = load volatile i32, ptr @result, align 4
+  %add = add nsw i32 %0, ptrtoint (ptr @func03_takeaddr to i32)
+  store volatile i32 %add, ptr @result, align 4
+  ret void
+}