[BOLT] Support computed goto and allow map addrs inside functions (#120267)

Ash Dobrescu · web-flow · commit 3bba26801368 · 2025-03-19T14:55:59.000Z
Create entry points for addresses referenced by dynamic relocations and allow getNewFunctionOrDataAddress to map addrs inside functions. By adding addresses referenced by dynamic relocations as entry points. This patch fixes an issue where bolt fails on code using computing goto's. This also fixes a mapping issue with the bugfix from this PR: #117766.
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -2455,6 +2455,30 @@ void RewriteInstance::readDynamicRelocations(const SectionRef &Section,
     if (Symbol)
       SymbolIndex[Symbol] = getRelocationSymbol(InputFile, Rel);
 
+    const uint64_t ReferencedAddress = SymbolAddress + Addend;
+    BinaryFunction *Func =
+        BC->getBinaryFunctionContainingAddress(ReferencedAddress);
+
+    if (Relocation::isRelative(RType) && SymbolAddress == 0) {
+      if (Func) {
+        if (!Func->isInConstantIsland(ReferencedAddress)) {
+          if (const uint64_t ReferenceOffset =
+                  ReferencedAddress - Func->getAddress()) {
+            Func->addEntryPointAtOffset(ReferenceOffset);
+          }
+        } else {
+          BC->errs() << "BOLT-ERROR: referenced address at 0x"
+                     << Twine::utohexstr(ReferencedAddress)
+                     << " is in constant island of function " << *Func << "\n";
+          exit(1);
+        }
+      }
+    } else if (Relocation::isRelative(RType) && SymbolAddress != 0) {
+      BC->errs() << "BOLT-ERROR: symbol address non zero for RELATIVE "
+                    "relocation type\n";
+      exit(1);
+    }
+
     BC->addDynamicRelocation(Rel.getOffset(), Symbol, RType, Addend);
   }
 }
@@ -5696,7 +5720,7 @@ uint64_t RewriteInstance::getNewFunctionOrDataAddress(uint64_t OldAddress) {
         for (const BinaryBasicBlock &BB : *BF)
           if (BB.isEntryPoint() &&
               (BF->getAddress() + BB.getOffset()) == OldAddress)
-            return BF->getOutputAddress() + BB.getOffset();
+            return BB.getOutputStartAddress();
       }
       BC->errs() << "BOLT-ERROR: unable to get new address corresponding to "
                     "input address 0x"
diff --git a/bolt/test/AArch64/computed-goto.s b/bolt/test/AArch64/computed-goto.s
@@ -0,0 +1,67 @@
+// This test checks that BOLT creates entry points for addresses
+// referenced by dynamic relocations.
+// The test also checks that BOLT can map addresses inside functions.
+
+// Checks for error and entry points.
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.bolt 2>&1 | FileCheck %s
+# RUN: llvm-bolt %t.exe -o %t.bolt --print-cfg | FileCheck --check-prefix=CHECK-ENTRIES %s
+
+// Checks for dynamic relocations.
+# RUN: llvm-readelf -dr %t.bolt > %t.out.txt
+# RUN: llvm-objdump -j .rela.dyn -d %t.bolt >> %t.out.txt
+# RUN: FileCheck --check-prefix=CHECK-RELOCS %s --input-file=%t.out.txt
+
+// Before bolt could handle mapping addresses within moved functions, it
+// would bail out with an error of the form:
+// BOLT-ERROR: unable to get new address corresponding to input address 0x10390 in function main. Consider adding this function to --skip-funcs=...
+// These addresses arise if computed GOTO is in use.
+// Check that bolt does not emit any error.
+# CHECK-NOT: BOLT-ERROR
+
+// Check that there are dynamic relocations.
+# CHECK-RELOCS:     Dynamic section at offset {{.*}} contains {{.*}} entries:
+# CHECK-RELOCS:     Relocation section '.rela.dyn' at offset {{.*}} contains {{.*}} entries
+
+// Check that dynamic relocations were updated
+# CHECK-RELOCS: [[#%x,OFF:]] [[#%x,INFO_DYN:]] R_AARCH64_RELATIVE [[#%x,ADDR:]]
+# CHECK-RELOCS-NEXT: [[#OFF + 8]] {{0*}}[[#INFO_DYN]] R_AARCH64_RELATIVE [[#ADDR + 8]]
+# CHECK-RELOCS: [[#ADDR]] <unknown>
+# CHECK-RELOCS: [[#ADDR + 8]] <unknown>
+
+// Check that BOLT registers extra entry points for dynamic relocations.
+# CHECK-ENTRIES: Binary Function "main" after building cfg {
+# CHECK-ENTRIES:  IsMultiEntry: 1
+# CHECK-ENTRIES: .Ltmp0 {{.*}}
+# CHECK-ENTRIES-NEXT: Secondary Entry Point: {{.*}}
+# CHECK-ENTRIES: .Ltmp1 {{.*}}
+# CHECK-ENTRIES-NEXT: Secondary Entry Point: {{.*}}
+
+.globl  main
+.p2align        2
+.type   main,@function
+main:
+.cfi_startproc
+        adrp    x8, .L__const.main.ptrs+8
+        add     x8, x8, :lo12:.L__const.main.ptrs+8
+        ldr     x9, [x8], #8
+        br      x9
+
+.Label0: // Block address taken
+        ldr     x9, [x8], #8
+        br      x9
+
+.Label1: // Block address taken
+        mov     w0, #42
+        ret
+
+.Lfunc_end0:
+.size   main, .Lfunc_end0-main
+.cfi_endproc
+        .type   .L__const.main.ptrs,@object
+        .section        .data.rel.ro,"aw",@progbits
+        .p2align        3, 0x0
+.L__const.main.ptrs:
+        .xword  .Label0
+        .xword  .Label1
diff --git a/bolt/test/Inputs/indirect_goto.c b/bolt/test/Inputs/indirect_goto.c
@@ -1,6 +1,6 @@
 int main(int argc, char *argv[]) {
-  static const void *T1[] = { &&L1, &&L2 };
-  static const void *T2[] = { &&L2, &&L3 };
+  static const void *T1[] = {&&L1, &&L2};
+  static const void *T2[] = {&&L2, &&L3};
 
   const void **T = (argc > 1) ? T1 : T2;
 
diff --git a/bolt/test/X86/indirect-goto-pie.test b/bolt/test/X86/indirect-goto-pie.test
diff --git a/bolt/test/X86/indirect-goto.test b/bolt/test/X86/indirect-goto.test
@@ -1,5 +1,5 @@
 ## Check llvm-bolt processes binaries compiled from sources that use indirect goto.
-RUN: %clang %cflags -no-pie %S/Inputs/indirect_goto.c -Wl,-q -o %t
+RUN: %clang %cflags -no-pie %S/../Inputs/indirect_goto.c -Wl,-q -o %t
 RUN: llvm-bolt %t -o %t.null --relocs=1 --print-cfg --print-only=main \
 RUN:   --strict \
 RUN:   2>&1 | FileCheck %s
diff --git a/bolt/test/indirect-goto-relocs.test b/bolt/test/indirect-goto-relocs.test
@@ -0,0 +1,19 @@
+// This test checks that BOLT creates entry points from sources
+// that use indirect goto.
+
+RUN: %clang %cflags -pie %S/Inputs/indirect_goto.c -o %t.exe -Wl,-q
+RUN: llvm-bolt %t.exe -o %t.bolt --print-cfg | FileCheck --check-prefix=CHECK-PIE %s
+
+RUN: %clang %cflags -no-pie %S/Inputs/indirect_goto.c -o %t.exe -Wl,-q
+RUN: llvm-bolt %t.exe -o %t.bolt --print-cfg | FileCheck --check-prefix=CHECK-NO-PIE %s
+
+// Check that BOLT registers extra entry points for dynamic relocations with PIE.
+CHECK-PIE: Binary Function "main" after building cfg {
+CHECK-PIE: IsMultiEntry: 1
+CHECK-PIE: Secondary Entry Points : {{.*}}
+
+// Check that BOLT does not register extra entry points for dynamic relocations
+// without PIE
+CHECK-NO-PIE: Binary Function "main" after building cfg {
+CHECK-NO-PIE-NOT: IsMultiEntry: 1
+CHECK-NO-PIE-NOT: Secondary Entry Points : {{.*}}
