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

[Vladislav Khmelevsky via llvm-commits]

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

>From f8cace02214774f20270bef30575fafa29cc3573 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 | 31 +++++++++++--------
 bolt/test/runtime/AArch64/controlflow.s       |  2 ++
 4 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/bolt/lib/Passes/ADRRelaxationPass.cpp b/bolt/lib/Passes/ADRRelaxationPass.cpp
index 7b612cbf6572bd3..4039fa2fbb51d48 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..fa9fb63c613dc13 100644
--- a/bolt/test/runtime/AArch64/adrrelaxationpass.s
+++ b/bolt/test/runtime/AArch64/adrrelaxationpass.s
@@ -1,33 +1,27 @@
 # 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
+# Also since main and test are 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
+# fail in non-strict mode
 
 # REQUIRES: system-linux
 
 # 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