From abd69b3653fa0298f717d535678a395f1cfa6bb4 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Tue, 27 Aug 2024 11:31:32 -0700
Subject: [PATCH] [BOLT] Handle internal calls in ValidateInternalCalls
 (#105736)

Move handling of all internal calls into the designated pass. Preserve
NOPs and mark functions as non-simple on non-X86 platforms.
---
 bolt/include/bolt/Core/BinaryFunction.h   |  2 +
 bolt/lib/Core/BinaryFunction.cpp          | 20 ++-----
 bolt/lib/Passes/ValidateInternalCalls.cpp |  9 ++--
 bolt/test/AArch64/internal-call.s         | 65 +++++++++++++++++++++++
 4 files changed, 77 insertions(+), 19 deletions(-)
 create mode 100644 bolt/test/AArch64/internal-call.s

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 24c7db2f5d69ca..6ebbaf94754e87 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -1692,6 +1692,8 @@ class BinaryFunction {
 
   void setPseudo(bool Pseudo) { IsPseudo = Pseudo; }
 
+  void setPreserveNops(bool Value) { PreserveNops = Value; }
+
   BinaryFunction &setUsesGnuArgsSize(bool Uses = true) {
     UsesGnuArgsSize = Uses;
     return *this;
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index af982fd60b17e7..46bdf208be6ad3 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1339,22 +1339,10 @@ Error BinaryFunction::disassemble() {
                   BC.getBinaryFunctionContainingAddress(TargetAddress))
             TargetFunc->setIgnored();
 
-        if (IsCall && containsAddress(TargetAddress)) {
-          if (TargetAddress == getAddress()) {
-            // Recursive call.
-            TargetSymbol = getSymbol();
-          } else {
-            if (BC.isX86()) {
-              // Dangerous old-style x86 PIC code. We may need to freeze this
-              // function, so preserve the function as is for now.
-              PreserveNops = true;
-            } else {
-              BC.errs() << "BOLT-WARNING: internal call detected at 0x"
-                        << Twine::utohexstr(AbsoluteInstrAddr)
-                        << " in function " << *this << ". Skipping.\n";
-              IsSimple = false;
-            }
-          }
+        if (IsCall && TargetAddress == getAddress()) {
+          // A recursive call. Calls to internal blocks are handled by
+          // ValidateInternalCalls pass.
+          TargetSymbol = getSymbol();
         }
 
         if (!TargetSymbol) {
diff --git a/bolt/lib/Passes/ValidateInternalCalls.cpp b/bolt/lib/Passes/ValidateInternalCalls.cpp
index 88df2e5b59f389..bdab895b6ac2e8 100644
--- a/bolt/lib/Passes/ValidateInternalCalls.cpp
+++ b/bolt/lib/Passes/ValidateInternalCalls.cpp
@@ -302,9 +302,6 @@ bool ValidateInternalCalls::analyzeFunction(BinaryFunction &Function) const {
 }
 
 Error ValidateInternalCalls::runOnFunctions(BinaryContext &BC) {
-  if (!BC.isX86())
-    return Error::success();
-
   // Look for functions that need validation. This should be pretty rare.
   std::set<BinaryFunction *> NeedsValidation;
   for (auto &BFI : BC.getBinaryFunctions()) {
@@ -312,14 +309,20 @@ Error ValidateInternalCalls::runOnFunctions(BinaryContext &BC) {
     for (BinaryBasicBlock &BB : Function) {
       for (MCInst &Inst : BB) {
         if (getInternalCallTarget(Function, Inst)) {
+          BC.errs() << "BOLT-WARNING: internal call detected in function "
+                    << Function << '\n';
           NeedsValidation.insert(&Function);
           Function.setSimple(false);
+          Function.setPreserveNops(true);
           break;
         }
       }
     }
   }
 
+  if (!BC.isX86())
+    return Error::success();
+
   // Skip validation for non-relocation mode
   if (!BC.HasRelocations)
     return Error::success();
diff --git a/bolt/test/AArch64/internal-call.s b/bolt/test/AArch64/internal-call.s
new file mode 100644
index 00000000000000..43b3a64f9d3280
--- /dev/null
+++ b/bolt/test/AArch64/internal-call.s
@@ -0,0 +1,65 @@
+## Test that llvm-bolt detects internal calls and marks the containing function
+## as non-simple.
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -static
+# RUN: llvm-bolt %t.exe -o %t.null --print-all 2>&1 | FileCheck %s
+
+# CHECK: Binary Function "_start" after building cfg
+# CHECK: internal call detected in function _start
+# CHECK-NOT: Binary Function "_start" after validate-internal-calls
+
+	.text
+  .globl _start
+  .type _start, %function
+_start:
+	.cfi_startproc
+.LBB00:
+	mov	x11, #0x1fff
+	cmp	x1, x11
+	b.hi	.Ltmp1
+
+.entry1:
+	movi	v4.16b, #0x0
+	movi	v5.16b, #0x0
+	subs	x1, x1, #0x8
+	b.lo	.Ltmp2
+
+.entry2:
+	ld1	{ v2.2d, v3.2d }, [x0], #32
+	ld1	{ v0.2d, v1.2d }, [x0], #32
+
+.Ltmp2:
+	uaddlp	v4.4s, v4.8h
+	uaddlp	v4.2d, v4.4s
+	mov	x0, v4.d[0]
+	mov	x1, v4.d[1]
+	add	x0, x0, x1
+	ret	x30
+
+.Ltmp1:
+	mov	x8, x30
+
+.Lloop:
+	add	x5, x0, x9
+	mov	x1, #0xface
+	movi	v4.16b, #0x0
+	movi	v5.16b, #0x0
+	bl	.entry2
+	add	x4, x4, x0
+	mov	x0, x5
+	sub	x7, x7, x10
+	cmp	x7, x11
+	b.hi	.Lloop
+
+	mov	x1, x7
+	bl	.entry1
+	add	x0, x4, x0
+	mov	x30, x8
+	ret	x30
+
+	.cfi_endproc
+.size _start, .-_start
+
+## Force relocation mode.
+  .reloc 0, R_AARCH64_NONE