Skip to content

[lld-macho] Use Symbols as branch target for safe_thunks ICF #126835

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 13, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 8 additions & 9 deletions lld/MachO/Arch/ARM64.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ struct ARM64 : ARM64Common {
void applyOptimizationHints(uint8_t *, const ObjFile &) const override;

void initICFSafeThunkBody(InputSection *thunk,
InputSection *branchTarget) const override;
InputSection *getThunkBranchTarget(InputSection *thunk) const override;
Symbol *targetSym) const override;
Symbol *getThunkBranchTarget(InputSection *thunk) const override;
uint32_t getICFSafeThunkSize() const override;
};

Expand Down Expand Up @@ -185,8 +185,7 @@ static constexpr uint32_t icfSafeThunkCode[] = {
0x14000000, // 08: b target
};

void ARM64::initICFSafeThunkBody(InputSection *thunk,
InputSection *branchTarget) const {
void ARM64::initICFSafeThunkBody(InputSection *thunk, Symbol *targetSym) const {
// The base data here will not be itself modified, we'll just be adding a
// reloc below. So we can directly use the constexpr above as the data.
thunk->data = {reinterpret_cast<const uint8_t *>(icfSafeThunkCode),
Expand All @@ -195,17 +194,17 @@ void ARM64::initICFSafeThunkBody(InputSection *thunk,
thunk->relocs.emplace_back(/*type=*/ARM64_RELOC_BRANCH26,
/*pcrel=*/true, /*length=*/2,
/*offset=*/0, /*addend=*/0,
/*referent=*/branchTarget);
/*referent=*/targetSym);
}

InputSection *ARM64::getThunkBranchTarget(InputSection *thunk) const {
Symbol *ARM64::getThunkBranchTarget(InputSection *thunk) const {
assert(thunk->relocs.size() == 1 &&
"expected a single reloc on ARM64 ICF thunk");
auto &reloc = thunk->relocs[0];
assert(isa<InputSection *>(reloc.referent) &&
"ARM64 thunk reloc is expected to point to an InputSection");
assert(isa<Symbol *>(reloc.referent) &&
"ARM64 thunk reloc is expected to point to a Symbol");

return cast<InputSection *>(reloc.referent);
return cast<Symbol *>(reloc.referent);
}

uint32_t ARM64::getICFSafeThunkSize() const { return sizeof(icfSafeThunkCode); }
Expand Down
49 changes: 37 additions & 12 deletions lld/MachO/ICF.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ using namespace lld;
using namespace lld::macho;

static constexpr bool verboseDiagnostics = false;
// This counter is used to generate unique thunk names.
static uint64_t icfThunkCounter = 0;

class ICF {
public:
Expand Down Expand Up @@ -263,6 +265,31 @@ void ICF::forEachClassRange(size_t begin, size_t end,
}
}

// Find or create a symbol at offset 0 in the given section
static Symbol *getThunkTargetSymbol(ConcatInputSection *isec) {
for (Symbol *sym : isec->symbols)
if (auto *d = dyn_cast<Defined>(sym))
if (d->value == 0)
return sym;

std::string thunkName;
if (isec->symbols.size() == 0)
thunkName = isec->getName().str() + ".icf.0";
else
thunkName = isec->getName().str() + "icf.thunk.target" +
std::to_string(icfThunkCounter++);

// If no symbol found at offset 0, create one
auto *sym = make<Defined>(thunkName, /*file=*/nullptr, isec,
/*value=*/0, /*size=*/isec->getSize(),
/*isWeakDef=*/false, /*isExternal=*/false,
/*isPrivateExtern=*/false, /*isThumb=*/false,
/*isReferencedDynamically=*/false,
/*noDeadStrip=*/false);
isec->symbols.push_back(sym);
return sym;
}

// Given a range of identical icfInputs, replace address significant functions
// with a thunk that is just a direct branch to the first function in the
// series. This way we keep only one main body of the function but we still
Expand All @@ -280,6 +307,9 @@ void ICF::applySafeThunksToRange(size_t begin, size_t end) {
// all thunks will branch to.
ConcatInputSection *masterIsec = icfInputs[begin];

// Get the symbol that all thunks will branch to.
Symbol *masterSym = getThunkTargetSymbol(masterIsec);

for (size_t i = begin + 1; i < end; ++i) {
ConcatInputSection *isec = icfInputs[i];
// When we're done processing keepUnique entries, we can stop. Sorting
Expand All @@ -291,7 +321,7 @@ void ICF::applySafeThunksToRange(size_t begin, size_t end) {
makeSyntheticInputSection(isec->getSegName(), isec->getName());
addInputSection(thunk);

target->initICFSafeThunkBody(thunk, masterIsec);
target->initICFSafeThunkBody(thunk, masterSym);
thunk->foldIdentical(isec, Symbol::ICFFoldKind::Thunk);

// Since we're folding the target function into a thunk, we need to adjust
Expand Down Expand Up @@ -495,18 +525,11 @@ Defined *macho::getBodyForThunkFoldedSym(Defined *foldedSym) {
// the actual body of the function.
InputSection *thunkBody = foldedSec->replacement;

// The actual (merged) body of the function that the thunk jumps to. This will
// end up in the final binary.
InputSection *functionBody = target->getThunkBranchTarget(thunkBody);

for (Symbol *sym : functionBody->symbols) {
Defined *d = dyn_cast<Defined>(sym);
// The symbol needs to be at the start of the InputSection
if (d && d->value == 0)
return d;
}
// The symbol of the merged body of the function that the thunk jumps to. This
// will end up in the final binary.
Symbol *targetSym = target->getThunkBranchTarget(thunkBody);

llvm_unreachable("could not find body symbol for ICF-generated thunk");
return cast<Defined>(targetSym);
}
void macho::foldIdenticalSections(bool onlyCfStrings) {
TimeTraceScope timeScope("Fold Identical Code Sections");
Expand All @@ -526,6 +549,8 @@ void macho::foldIdenticalSections(bool onlyCfStrings) {
// ICF::segregate()
std::vector<ConcatInputSection *> foldable;
uint64_t icfUniqueID = inputSections.size();
// Reset the thunk counter for each run of ICF.
icfThunkCounter = 0;
for (ConcatInputSection *isec : inputSections) {
bool isFoldableWithAddendsRemoved = isCfStringSection(isec) ||
isClassRefsSection(isec) ||
Expand Down
4 changes: 2 additions & 2 deletions lld/MachO/Target.h
Original file line number Diff line number Diff line change
Expand Up @@ -76,13 +76,13 @@ class TargetInfo {

// Init 'thunk' so that it be a direct jump to 'branchTarget'.
virtual void initICFSafeThunkBody(InputSection *thunk,
InputSection *branchTarget) const {
Symbol *targetSym) const {
llvm_unreachable("target does not support ICF safe thunks");
}

// Given a thunk for which `initICFSafeThunkBody` was called, return the
// branchTarget it was initialized with.
virtual InputSection *getThunkBranchTarget(InputSection *thunk) const {
virtual Symbol *getThunkBranchTarget(InputSection *thunk) const {
llvm_unreachable("target does not support ICF safe thunks");
}

Expand Down
35 changes: 33 additions & 2 deletions lld/test/MachO/arm64-thunks.s
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,16 @@

# RUN: rm -rf %t; mkdir %t
# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t/input.o
# RUN: %lld -arch arm64 -dead_strip -lSystem -U _extern_sym -map %t/thunk.map -o %t/thunk %t/input.o
## Use --icf=safe_thunks to test that branch extension algo is compatible
## with safe_thunks ICF.
# RUN: %lld -arch arm64 -dead_strip -lSystem -U _extern_sym -map %t/thunk.map -o %t/thunk %t/input.o --icf=safe_thunks
# RUN: llvm-objdump --no-print-imm-hex -d --no-show-raw-insn %t/thunk | FileCheck %s

# RUN: FileCheck %s --input-file %t/thunk.map --check-prefix=MAP

# MAP: 0x{{[[:xdigit:]]+}} {{.*}} _b
# MAP: 0x{{[[:xdigit:]]+}} {{.*}} _fold_func_low_addr
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _a
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _b
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _c
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _d.thunk.0
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _e.thunk.0
Expand All @@ -35,10 +39,13 @@
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _b.thunk.0
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _h
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _main
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _fold_func_high_addr
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _c.thunk.0
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _d.thunk.1
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _e.thunk.1
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _f.thunk.1
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _fold_func_low_addr.thunk.0
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} ltmp0.thunk.0
# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _z


Expand Down Expand Up @@ -200,8 +207,21 @@
# CHECK: [[#%x, NAN_PAGE + NAN_OFFSET]] <__stubs>:

.subsections_via_symbols

.addrsig
.addrsig_sym _fold_func_low_addr
.addrsig_sym _fold_func_high_addr

.text

.globl _fold_func_low_addr
.p2align 2
_fold_func_low_addr:
add x0, x0, x0
add x1, x0, x1
add x2, x0, x2
ret

.globl _a
.p2align 2
_a:
Expand Down Expand Up @@ -329,9 +349,20 @@ _main:
bl _f
bl _g
bl _h
bl _fold_func_low_addr
bl _fold_func_high_addr
bl ___nan
ret

.globl _fold_func_high_addr
.p2align 2
_fold_func_high_addr:
add x0, x0, x0
add x1, x0, x1
add x2, x0, x2
ret


.section __TEXT,__cstring
# The .space below has to be composed of non-zero characters. Otherwise, the
# linker will create a symbol for every '0' in the section, leading to
Expand Down