[llvm] [BOLT][AArch64] Fix strict usage during ADR Relax (PR #71377)

[Vladislav Khmelevsky via llvm-commits]

https://github.com/yota9 created https://github.com/llvm/llvm-project/pull/71377

Currently strict mode is used to expand number of optimized functions,
not to shrink it. Revert the option usage in the pass, so passing strict
option would relax adr instruction even if there are no nops around it.
Also add check for nop after adr instruction.


>From ebc7188877360245c1bd638d8074eefc3b06777e Mon Sep 17 00:00:00 2001
From: Vladislav Khmelevsky <och95 at yandex.ru>
Date: Mon, 6 Nov 2023 12:36:57 +0400
Subject: [PATCH] [BOLT][AArch64] Fix strict usage during ADR Relax

Currently strict mode is used to expand number of optimized functions,
not to shrink it. Revert the option usage in the pass, so passing strict
option would relax adr instruction even if there are no nops around it.
Also add check for nop after adr instruction.
---
 bolt/lib/Passes/ADRRelaxationPass.cpp         |  7 ++++--
 bolt/test/AArch64/r_aarch64_prelxx.s          |  2 +-
 bolt/test/runtime/AArch64/adrrelaxationpass.s | 25 +++++++++++--------
 bolt/test/runtime/AArch64/controlflow.s       |  2 ++
 4 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/bolt/lib/Passes/ADRRelaxationPass.cpp b/bolt/lib/Passes/ADRRelaxationPass.cpp
index 7b612cbf6572bd3..c5cfb4901a9ba5d 100644
--- a/bolt/lib/Passes/ADRRelaxationPass.cpp
+++ b/bolt/lib/Passes/ADRRelaxationPass.cpp
@@ -72,14 +72,17 @@ void ADRRelaxationPass::runOnFunction(BinaryFunction &BF) {
 
       if (It != BB.begin() && BC.MIB->isNoop(*std::prev(It))) {
         It = BB.eraseInstruction(std::prev(It));
-      } else if (opts::StrictMode && !BF.isSimple()) {
+      } else if (std::next(It) != BB.end() && BC.MIB->isNoop(*std::next(It))) {
+          BB.eraseInstruction(std::next(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.
+        auto L = BC.scopeLock();
         errs() << formatv("BOLT-ERROR: Cannot relax adr in non-simple function "
-                          "{0}. Can't proceed in current mode.\n",
+                          "{0}. Use --strict option to override\n",
                           BF.getOneName());
         PassFailed = true;
         return;
diff --git a/bolt/test/AArch64/r_aarch64_prelxx.s b/bolt/test/AArch64/r_aarch64_prelxx.s
index 444dee72b7c04eb..73bf8387d363451 100644
--- a/bolt/test/AArch64/r_aarch64_prelxx.s
+++ b/bolt/test/AArch64/r_aarch64_prelxx.s
@@ -12,7 +12,7 @@
 // CHECKPREL-NEXT:  R_AARCH64_PREL32      {{.*}} _start + 4
 // CHECKPREL-NEXT:  R_AARCH64_PREL64      {{.*}} _start + 8
 
-// RUN: llvm-bolt %t.exe -o %t.bolt
+// RUN: llvm-bolt %t.exe -o %t.bolt --strict
 // RUN: llvm-objdump -D %t.bolt | FileCheck %s --check-prefix=CHECKPREL32
 
 // CHECKPREL32: [[#%x,DATATABLEADDR:]] <datatable>:
diff --git a/bolt/test/runtime/AArch64/adrrelaxationpass.s b/bolt/test/runtime/AArch64/adrrelaxationpass.s
index 5c50cd6371926cb..12482d04d60430c 100644
--- a/bolt/test/runtime/AArch64/adrrelaxationpass.s
+++ b/bolt/test/runtime/AArch64/adrrelaxationpass.s
@@ -10,24 +10,18 @@
 # RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
 # RUN:   %s -o %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
-# RUN: llvm-bolt %t.exe -o %t.bolt --adr-relaxation=true
+# RUN: llvm-bolt %t.exe -o %t.bolt --adr-relaxation=true --strict
 # 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: not llvm-bolt %t.exe -o %t.bolt --adr-relaxation=true \
 # RUN: 2>&1 | FileCheck %s --check-prefix CHECK-ERROR
 
-  .data
-  .align 8
-  .global Gvar
-Gvar: .xword 0x0
-  .global Gvar2
-Gvar2: .xword 0x42
-
   .text
   .align 4
   .global test
   .type test, %function
 test:
+  adr x2, Gvar
   mov x0, xzr
   ret
   .size test, .-test
@@ -47,6 +41,17 @@ br:
 .CI:
   .word 0xff
 
+  .data
+  .align 8
+  .global Gvar
+Gvar: .xword 0x0
+  .global Gvar2
+Gvar2: .xword 0x42
+  .balign 4
+jmptable:
+  .word 0
+  .word test - jmptable
+
 # CHECK: <main>:
 # CHECK-NEXT: adr x0, 0x{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: adrp x1, 0x{{[1-8a-f][0-9a-f]*}}
@@ -54,4 +59,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, 0x{{[1-8a-f][0-9a-f]*}}
-# CHECK-ERROR: BOLT-ERROR: Cannot relax adr in non-simple function main
+# CHECK-ERROR: BOLT-ERROR: Cannot relax adr in non-simple function
diff --git a/bolt/test/runtime/AArch64/controlflow.s b/bolt/test/runtime/AArch64/controlflow.s
index fe9aab88f0c7404..7b0a38779f6e9cf 100644
--- a/bolt/test/runtime/AArch64/controlflow.s
+++ b/bolt/test/runtime/AArch64/controlflow.s
@@ -48,6 +48,7 @@ test_cond_branch:
   .global test_branch_reg
   .type test_branch_reg, %function
 test_branch_reg:
+  nop
   adr x0, test_branch_zero
   br x0
   panic
@@ -97,6 +98,7 @@ test_call:
   .global test_call_reg
   .type test_call_reg, %function
 test_call_reg:
+  nop
   adr x0, test_call_foo
   blr x0
   panic