[ellishg]

With --icf=safe_thunks, address significant functions transform into a thunk that branches to the folded function body. Those thunks are necessary for instructions using the address and indirect branches, but a direct branch to the thunk can be safely redirected to the function body.

This PR allows folding functions if one has a direct call to a function body and another has a direct call to the thunk (and are otherwise identical).

This allows us to ICF more functions and reduce the number of direct calls to thunks.

LLM assisted

[llvmorg-github-actions]

@llvm/pr-subscribers-lld-macho

@llvm/pr-subscribers-lld

Author: Ellis Hoag (ellishg)

Changes

With --icf=safe_thunks, address significant functions transform into a thunk that branches to the folded function body. Those thunks are necessary for instructions using the address and indirect branches, but a direct branch to the thunk can be safely redirected to the function body.

This PR allows folding functions if one has a direct call to a function body and another has a direct call to the thunk (and are otherwise identical).

This allows us to ICF more functions and reduce the number of direct calls to thunks.

LLM assisted

---

Full diff: https://github.com/llvm/llvm-project/pull/239470.diff

3 Files Affected:

• (modified) lld/MachO/ICF.cpp (+19-3)
• (added) lld/test/MachO/icf-safe-thunks-callers.s (+219)
• (modified) lld/test/MachO/icf-safe-thunks.ll (+4-8)

diff --git a/lld/MachO/ICF.cpp b/lld/MachO/ICF.cpp
index 9169463792bd7..0d49a9fa7639f 100644
--- a/lld/MachO/ICF.cpp
+++ b/lld/MachO/ICF.cpp
@@ -142,6 +142,17 @@ static bool compareData(const ConcatInputSection *ia,
   return bufA == bufB;
 }
 
+static Defined *getDirectCallTarget(const Relocation &r) {
+  if (r.getAddend() || !target->hasAttr(r.type, RelocAttrBits::BRANCH))
+    return nullptr;
+  auto *sym = dyn_cast_if_present<Symbol *>(r.referent);
+  auto *d = dyn_cast_if_present<Defined>(sym);
+  if (!d || d->value || !d->originalIsec || !isCodeSection(d->originalIsec) ||
+      needsBinding(sym))
+    return nullptr;
+  return d;
+}
+
 // Compare "non-moving" parts of two ConcatInputSections, namely everything
 // except references to other ConcatInputSections.
 bool ICF::equalsConstant(const ConcatInputSection *ia,
@@ -204,9 +215,14 @@ bool ICF::equalsConstant(const ConcatInputSection *ia,
     // relocations to reference identical `keepUnique` functions that will be
     // distinguished later via thunks - so we need to handle this case
     // explicitly.
-    if ((isecA != isecB) && ((isecA->keepUnique && isCodeSection(isecA)) ||
-                             (isecB->keepUnique && isCodeSection(isecB))))
-      return false;
+    if (isecA != isecB && ((isecA->keepUnique && isCodeSection(isecA)) ||
+                           (isecB->keepUnique && isCodeSection(isecB)))) {
+      // In `safe_thunks` mode, we can still fold functions with direct calls to
+      // equivalent functions.
+      if (config->icfLevel != ICFLevel::safe_thunks ||
+          !getDirectCallTarget(ra) || !getDirectCallTarget(rb))
+        return false;
+    }
 
     if (isecA->parent != isecB->parent)
       return false;
diff --git a/lld/test/MachO/icf-safe-thunks-callers.s b/lld/test/MachO/icf-safe-thunks-callers.s
new file mode 100644
index 0000000000000..9add5c0806bc9
--- /dev/null
+++ b/lld/test/MachO/icf-safe-thunks-callers.s
@@ -0,0 +1,219 @@
+# REQUIRES: aarch64
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t.o
+# RUN: %lld -arch arm64 -dylib --icf=safe_thunks -interposable %t.o -map - -o %t.dylib | FileCheck %s
+# RUN: llvm-objdump --macho --disassemble %t.dylib | FileCheck %s --check-prefix=DISASM
+
+## Taking the leaves' addresses must still produce different code. A caller of
+## a thunk leading to a different body must also keep its own body.
+# CHECK-DAG: 0x0000000C {{.*}} _address_a
+# CHECK-DAG: 0x0000000C {{.*}} _address_b
+# CHECK-DAG: 0x00000008 {{.*}} _leaf_c
+# CHECK-DAG: 0x00000004 {{.*}} _leaf_d
+# CHECK-DAG: 0x00000004 {{.*}} _interposable_leaf
+# CHECK-DAG: 0x00000018 {{.*}} _call_c
+
+## Calls and tail branches to equivalent bodies allow their callers to fold.
+# CHECK-DAG: [[CALL:0x[0-9A-Fa-f]+]] 0x00000018 {{.*}} _call_a
+# CHECK-DAG: [[CALL]] 0x00000000 {{.*}} _call_b
+# CHECK-DAG: [[TAIL:0x[0-9A-Fa-f]+]] 0x00000008 {{.*}} _tail_a
+# CHECK-DAG: [[TAIL]] 0x00000000 {{.*}} _tail_b
+
+## Address-significant callers must retain unique entrypoints too.
+# CHECK-DAG: 0x00000018 {{.*}} _unique_a
+# CHECK-DAG: 0x00000004 {{.*}} _unique_b
+
+## Recursive bodies can be equivalent while their entrypoints remain distinct.
+# CHECK-DAG: 0x00000010 {{.*}} _recursive_a
+# CHECK-DAG: 0x00000004 {{.*}} _recursive_b
+
+## Address-significant leaves share a body through a unique entry thunk.
+# DISASM-LABEL: _leaf_a:
+# DISASM-NEXT: [[#%x,LEAF_A:]]:
+
+## Folded callers share a callsite; distinct callees keep distinct callers.
+# DISASM-LABEL: _call_b:
+# DISASM: {{.*}} bl _leaf_a
+# DISASM-LABEL: _call_c:
+# DISASM: {{.*}} bl _leaf_d
+
+# DISASM-LABEL: _call_only_b:
+# DISASM: {{.*}} bl _leaf_b
+
+# DISASM-LABEL: _address_a:
+# DISASM: add x0, x0, #0x[[#%x,LEAF_A]]
+# DISASM-LABEL: _address_b:
+# DISASM: add x0, x0, #0x[[#%x,LEAF_B:]]
+## Direct and indirect calls to a thunk preserve the same entrypoint.
+# DISASM-LABEL: _call_direct_and_indirect_b:
+# DISASM: add x8, x8, #0x[[#%x,LEAF_B]]
+# DISASM: {{.*}} bl _leaf_b
+# DISASM: blr x8
+
+## Runtime-bound direct calls still resolve the original symbol.
+# DISASM-LABEL: _call_interposable:
+# DISASM: {{.*}} bl {{.*}}symbol stub for: _interposable_leaf
+
+# DISASM-LABEL: _leaf_d:
+# DISASM-NEXT: {{.*}} b _leaf_c
+# DISASM-LABEL: _leaf_b:
+# DISASM-NEXT: [[#%x,LEAF_B]]: {{.*}} b _leaf_a
+# DISASM-LABEL: _unique_b:
+# DISASM-NEXT: {{.*}} b _unique_a
+
+.subsections_via_symbols
+.text
+.p2align 2
+.globl _leaf_a, _leaf_b, _leaf_c, _leaf_d
+.globl _call_a, _call_b, _call_c, _call_only_b, _tail_a, _tail_b
+.globl _unique_a, _unique_b, _address_a, _address_b
+.globl _recursive_a, _recursive_b
+.globl _call_direct_and_indirect_b
+.private_extern _leaf_a, _leaf_b, _leaf_c, _leaf_d
+.private_extern _unique_a, _unique_b, _recursive_a, _recursive_b
+.globl _interposable_leaf, _call_interposable
+
+_leaf_a:
+  add w0, w0, #7
+  ret
+
+_leaf_b:
+  add w0, w0, #7
+  ret
+
+_leaf_c:
+  add w0, w0, #9
+  ret
+
+_leaf_d:
+  add w0, w0, #9
+  ret
+
+_call_a:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_a
+  add w0, w0, #3
+  ldp x29, x30, [sp], #16
+  ret
+
+_call_b:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_b
+  add w0, w0, #3
+  ldp x29, x30, [sp], #16
+  ret
+
+_call_c:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_d
+  add w0, w0, #3
+  ldp x29, x30, [sp], #16
+  ret
+
+_call_only_b:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_b
+  add w0, w0, #11
+  ldp x29, x30, [sp], #16
+  ret
+
+_tail_a:
+  eor w0, w0, #1
+  b _leaf_a
+
+_tail_b:
+  eor w0, w0, #1
+  b _leaf_b
+
+_unique_a:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_a
+  add w0, w0, #5
+  ldp x29, x30, [sp], #16
+  ret
+
+_unique_b:
+  stp x29, x30, [sp, #-16]!
+  mov x29, sp
+  bl _leaf_b
+  add w0, w0, #5
+  ldp x29, x30, [sp], #16
+  ret
+
+_recursive_a:
+  subs w0, w0, #1
+  b.eq Lrecursive_done_a
+  b _recursive_step_a
+Lrecursive_done_a:
+  ret
+
+_recursive_b:
+  subs w0, w0, #1
+  b.eq Lrecursive_done_b
+  b _recursive_step_b
+Lrecursive_done_b:
+  ret
+
+_recursive_step_a:
+  b _recursive_a
+  nop
+
+_recursive_step_b:
+  b _recursive_b
+  nop
+
+_address_a:
+  adrp x0, _leaf_a@PAGE
+  add x0, x0, _leaf_a@PAGEOFF
+  ret
+
+_address_b:
+  adrp x0, _leaf_b@PAGE
+  add x0, x0, _leaf_b@PAGEOFF
+  ret
+
+_call_direct_and_indirect_b:
+  stp x29, x30, [sp, #-32]!
+  mov x29, sp
+  adrp x8, _leaf_b@PAGE
+  add x8, x8, _leaf_b@PAGEOFF
+  str x8, [sp, #16]
+  bl _leaf_b
+  ldr x8, [sp, #16]
+  blr x8
+  ldp x29, x30, [sp], #32
+  ret
+
+_interposable_leaf:
+  add w0, w0, #7
+  ret
+
+_call_interposable:
+  stp x29, x30, [sp, #-16]!
+  bl _interposable_leaf
+  ldp x29, x30, [sp], #16
+  ret
+
+.data
+.p2align 3
+.globl _addresses
+_addresses:
+  .quad _leaf_a, _leaf_b, _leaf_c, _leaf_d, _unique_a, _unique_b, _recursive_a, _recursive_b
+
+.addrsig
+.addrsig_sym _leaf_a
+.addrsig_sym _leaf_b
+.addrsig_sym _leaf_c
+.addrsig_sym _leaf_d
+.addrsig_sym _unique_a
+.addrsig_sym _unique_b
+.addrsig_sym _recursive_a
+.addrsig_sym _recursive_b
+.addrsig_sym _recursive_step_a
+.addrsig_sym _recursive_step_b
+.addrsig_sym _interposable_leaf
diff --git a/lld/test/MachO/icf-safe-thunks.ll b/lld/test/MachO/icf-safe-thunks.ll
index 1d92fd81d6377..0953f6e8a4b88 100644
--- a/lld/test/MachO/icf-safe-thunks.ll
+++ b/lld/test/MachO/icf-safe-thunks.ll
@@ -32,9 +32,7 @@
 ; CHECK-ARM64:             mov {{.*}}, #0x21
 ;
 ; CHECK-ARM64:        _func_call_thunked_1_nomerge:
-; CHECK-ARM64-NEXT:        stp	x29
-;
-; CHECK-ARM64:        _func_call_thunked_2_nomerge:
+; CHECK-ARM64-NEXT:   _func_call_thunked_2_nomerge:
 ; CHECK-ARM64-NEXT:   _func_call_thunked_2_merge:
 ; CHECK-ARM64-NEXT:        stp	x29
 ;
@@ -69,9 +67,7 @@
 ; CHECK-LTO:             mov {{.*}}, #0x21
 ;
 ; CHECK-LTO:        _func_call_thunked_1_nomerge:
-; CHECK-LTO-NEXT:        stp	x29
-;
-; CHECK-LTO:        _func_call_thunked_2_nomerge:
+; CHECK-LTO-NEXT:   _func_call_thunked_2_nomerge:
 ; CHECK-LTO-NEXT:   _func_call_thunked_2_merge:
 ; CHECK-LTO-NEXT:        stp	x29
 ;
@@ -96,7 +92,7 @@
 ; CHECK-LTO-MAP-NEXT: 0x00000000 [  2] _func_3identical_v2_canmerge
 ; CHECK-LTO-MAP-NEXT: 0x00000000 [  2] _func_3identical_v3_canmerge
 ; CHECK-LTO-MAP-NEXT: 0x00000020 [  2] _func_call_thunked_1_nomerge
-; CHECK-LTO-MAP-NEXT: 0x00000020 [  2] _func_call_thunked_2_nomerge
+; CHECK-LTO-MAP-NEXT: 0x00000000 [  2] _func_call_thunked_2_nomerge
 ; CHECK-LTO-MAP-NEXT: 0x00000000 [  2] _func_call_thunked_2_merge
 ; CHECK-LTO-MAP-NEXT: 0x00000034 [  2] _call_all_funcs
 ; CHECK-LTO-MAP-NEXT: 0x00000050 [  2] _take_func_addr
@@ -112,7 +108,7 @@
 ; CHECK-ARM64-MAP-NEXT: 0x00000000 [  2] _func_3identical_v2_canmerge
 ; CHECK-ARM64-MAP-NEXT: 0x00000000 [  2] _func_3identical_v3_canmerge
 ; CHECK-ARM64-MAP-NEXT: 0x00000020 [  2] _func_call_thunked_1_nomerge
-; CHECK-ARM64-MAP-NEXT: 0x00000020 [  2] _func_call_thunked_2_nomerge
+; CHECK-ARM64-MAP-NEXT: 0x00000000 [  2] _func_call_thunked_2_nomerge
 ; CHECK-ARM64-MAP-NEXT: 0x00000000 [  2] _func_call_thunked_2_merge
 ; CHECK-ARM64-MAP-NEXT: 0x00000034 [  2] _call_all_funcs
 ; CHECK-ARM64-MAP-NEXT: 0x00000050 [  2] _take_func_addr

[nico]

Sorry, I accidentally clicked "Request" next to Copilot while trying to add Andy. That wasn't intentional. I can't see how to cancel.

[nico]

@andysjlim, want to take a look? (apparently I can't add you as a formal reviewer; probably because you're not yet an llvm commiter)

[Copilot]

🟡 Changes recommended

The new test’s differing instruction counts cause compareData() to reject the functions before reaching needsBinding(), leaving the interposition safeguard untested until the instruction sequences are made otherwise identical.

1 open issue

• Make interposable-call test reach the callee comparison · New

What changed in this PR

This PR improves Mach-O identical code folding under --icf=safe_thunks, allowing otherwise identical callers of equivalent function bodies and unique entry thunks to fold together.

Changes:

• Allow eligible direct branch targets to participate in caller folding while excluding runtime-bound symbols.
• Add coverage for calls, tail branches, recursion, address identity, and interposition.
• Update existing ARM64 and LTO expectations for additional folding.

File	Description
lld/​test/​MachO/​icf-safe-thunks.ll	Update disassembly and map expectations for folded callers.
lld/​test/​MachO/​icf-safe-thunks-callers.s	Add focused safe-thunk caller-folding regression tests.
lld/​MachO/​ICF.cpp	Permit folding through eligible direct calls to equivalent address-significant functions.

_{Review effort: Balanced}