[BOLT][AArch64] Replace NOP with adrp in AdrRelaxationPass to preserve relative offsets.

Avoid replacing one adr instruction with two adrp+add by utilizing linker-provided nops
when they are present. By doing so we preserve relative offsets of next instructions
in a function which reduces chances to break undetected jump tables. This commit makes
release-mode lld-linked clang, lld and etc work after BOLT.

Reviewed By: rafauler, yota9

Differential Revision: https://reviews.llvm.org/D143887

2 files changed, 24 insertions, 0 deletions

diff --git a/bolt/lib/Passes/ADRRelaxationPass.cpp b/bolt/lib/Passes/ADRRelaxationPass.cpp
index 4d593678c2e1..76924d96fcf9 100644
--- a/bolt/lib/Passes/ADRRelaxationPass.cpp
+++ b/bolt/lib/Passes/ADRRelaxationPass.cpp
@@ -12,6 +12,8 @@
 
 #include "bolt/Passes/ADRRelaxationPass.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Utils/CommandLineOpts.h"
+#include <iterator>
 
 using namespace llvm;
 
@@ -54,6 +56,20 @@ void ADRRelaxationPass::runOnFunction(BinaryFunction &BF) {
       int64_t Addend = BC.MIB->getTargetAddend(Inst);
       InstructionListType Addr =
           BC.MIB->materializeAddress(Symbol, BC.Ctx.get(), Reg, Addend);
+
+      if (It != BB.begin() && BC.MIB->isNoop(*std::prev(It))) {
+        It = BB.eraseInstruction(std::prev(It));
+      } else if (opts::StrictMode && !BF.isSimple()) {
+        // If the function is not simple, it may contain a jump table undetected
+        // by us. This jump table may use an offset from the branch instruction
+        // to land in the desired place. If we add new instructions, we
+        // invalidate this offset, so we have to rely on linker-inserted NOP to
+        // replace it with ADRP, and abort if it is not present.
+        errs() << formatv("BOLT-ERROR: Cannot relax adr in non-simple function "
+                          "{0}. Can't proceed in current mode.\n",
+                          BF.getOneName());
+        exit(1);
+      }
       It = BB.replaceInstruction(It, Addr);
     }
   }
diff --git a/bolt/test/runtime/AArch64/adrrelaxationpass.s b/bolt/test/runtime/AArch64/adrrelaxationpass.s
index 4c93dfe19c31..245c60979b3b 100644
--- a/bolt/test/runtime/AArch64/adrrelaxationpass.s
+++ b/bolt/test/runtime/AArch64/adrrelaxationpass.s
@@ -1,5 +1,9 @@
 # The second and third ADR instructions are non-local to functions
 # and must be replaced with ADRP + ADD by BOLT
+# Also since main is non-simple, we can't change it's length so we have to
+# replace NOP with adrp, and if there is no nop before adr in non-simple
+# function, we can't guarantee we didn't break possible jump tables, so we
+# fail in strict mode
 
 # REQUIRES: system-linux
 
@@ -9,6 +13,8 @@
 # RUN: llvm-bolt %t.exe -o %t.bolt --adr-relaxation=true
 # RUN: llvm-objdump --no-print-imm-hex -d --disassemble-symbols=main %t.bolt | FileCheck %s
 # RUN: %t.bolt
+# RUN: not llvm-bolt %t.exe -o %t.bolt --adr-relaxation=true --strict \
+# RUN: 2>&1 | FileCheck %s --check-prefix CHECK-ERROR
 
   .data
   .align 8
@@ -31,6 +37,7 @@ test:
   .type main, %function
 main:
   adr x0, .CI
+  nop
   adr x1, test
   adr x2, Gvar2
   adr x3, br
@@ -47,3 +54,4 @@ br:
 # CHECK-NEXT: adrp x2, 0x{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: add x2, x2, #{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: adr x3, #{{[0-9][0-9]*}}
+# CHECK-ERROR: BOLT-ERROR: Cannot relax adr in non-simple function main