From 10d9659bdf9a908907f18603820abd8365db1b4a Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Mon, 27 Jan 2025 22:14:33 -0800
Subject: [PATCH] [BOLT][AArch64] Remove nops in functions with defined control
 flow

When a function has an indirect branch with unknown control flow, we
preserve nops in order to keep all instruction offsets (from the start
of the function) the same in case the indirect branch is used by a
PC-relative jump table. However, when we know the control flow of the
function, we should be able to safely remove nops.
---
 bolt/lib/Core/BinaryFunction.cpp |  5 ++++-
 bolt/test/AArch64/remove-nops.s  | 28 ++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 bolt/test/AArch64/remove-nops.s

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 1c5cd62a095b2..88562a60dd50d 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -795,7 +795,6 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
 
   auto Begin = Instructions.begin();
   if (BC.isAArch64()) {
-    PreserveNops = BC.HasRelocations;
     // Start at the last label as an approximation of the current basic block.
     // This is a heuristic, since the full set of labels have yet to be
     // determined
@@ -2300,6 +2299,10 @@ Error BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
       BC.errs() << "BOLT-WARNING: failed to post-process indirect branches for "
                 << *this << '\n';
     }
+
+    if (BC.isAArch64())
+      PreserveNops = BC.HasRelocations;
+
     // In relocation mode we want to keep processing the function but avoid
     // optimizing it.
     setSimple(false);
diff --git a/bolt/test/AArch64/remove-nops.s b/bolt/test/AArch64/remove-nops.s
new file mode 100644
index 0000000000000..0f02a4b273dda
--- /dev/null
+++ b/bolt/test/AArch64/remove-nops.s
@@ -0,0 +1,28 @@
+## Verify that llvm-bolt removes nop instructions from functions with indirect
+## branches that have defined control flow.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags --target=aarch64-unknown-linux %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.bolt --print-normalized 2>&1 | FileCheck %s
+# RUN: llvm-objdump -d --disassemble-symbols=_start %t.bolt \
+# RUN:   | FileCheck %s --check-prefix=CHECK-OBJDUMP
+
+# CHECK-OBJDUMP-LABEL: _start
+# CHECK-OBJDUMP-NOT: nop
+
+  .section .text
+  .align 4
+  .globl _start
+  .type  _start, %function
+_start:
+# CHECK-LABEL: Binary Function "_start"
+  nop
+# CHECK-NOT: nop
+  br      x0
+# CHECK: br x0 # TAILCALL
+.size _start, .-_start
+
+## Force relocation mode.
+  .reloc 0, R_AARCH64_NONE