[AArch64][v8.5A] Add BTI to all function starts

The existing BTI placement pass avoids inserting "BTI c" when the
function has local linkage and is only directly called. However,
even in this case, there is a (small) chance that the linker later
adds a hunk with an indirect call to the function, e.g. if the
function is placed in a separate section and moved far away from
its callers. Make sure to add BTI for these functions too.

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

Author: pbarrio

llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

@@ -64,7 +64,6 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(
       dbgs() << "********** AArch64 Branch Targets  **********\n"
              << "********** Function: " << MF.getName() << '\n');
-  const Function &F = MF.getFunction();
 
   // LLVM does not consider basic blocks which are the targets of jump tables
   // to be address-taken (the address can't escape anywhere else), but they are
@@ -78,13 +77,16 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   bool MadeChange = false;
   for (MachineBasicBlock &MBB : MF) {
     bool CouldCall = false, CouldJump = false;
-    // If the function is address-taken or externally-visible, it could be
-    // indirectly called. PLT entries and tail-calls use BR, but when they are
+    // Even in cases where a function has internal linkage and is only called
+    // directly in its translation unit, it can still be called indirectly if
+    // the linker decides to add a thunk to it for whatever reason (say, for
+    // example, if it is finally placed far from its call site and a BL is not
+    // long-range enough). PLT entries and tail-calls use BR, but when they are
     // are in guarded pages should all use x16 or x17 to hold the called
     // address, so we don't need to set CouldJump here. BR instructions in
     // non-guarded pages (which might be non-BTI-aware code) are allowed to
     // branch to a "BTI c" using any register.
-    if (&MBB == &*MF.begin() && (F.hasAddressTaken() || !F.hasLocalLinkage()))
+    if (&MBB == &*MF.begin())
       CouldCall = true;
 
     // If the block itself is address-taken, it could be indirectly branched

llvm/test/CodeGen/AArch64/branch-target-enforcement.mir

@@ -116,12 +116,15 @@ body:             |
     RET undef $lr, implicit killed $w0
 
 ---
-# Internal function, not address-taken in this module, so no BTI needed.
+# Internal function, not address-taken in this module, however the compiler
+# cannot 100% ensure that later parts of the toolchain won't add indirect
+# jumps. E.g. a linker adding a thunk to extend the range of a direct jump.
+# Therefore, even this case needs a BTI.
 name:            simple_internal
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: simple_internal
-    ; CHECK-NOT: HINT
+    ; CHECK: HINT 34
     ; CHECK: RET
     $w0 = ORRWrs $wzr, $wzr, 0
     RET undef $lr, implicit killed $w0

llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll

@@ -48,19 +48,20 @@ define void @f2_1() "patchable-function-entry"="1" "patchable-function-prefix"="
 }
 
 ;; -fpatchable-function-entry=1 -mbranch-protection=bti
-;; For M=0, don't create .Lpatch0 if the initial instruction is not BTI,
-;; even if other basic blocks may have BTI.
+;; We add BTI c even when the function has internal linkage
 define internal void @f1i(i64 %v) "patchable-function-entry"="1" "branch-target-enforcement"="true" {
 ; CHECK-LABEL: f1i:
 ; CHECK-NEXT: .Lfunc_begin3:
 ; CHECK:      // %bb.0:
+; CHECK-NEXT:  hint #34
+; CHECK-NEXT:  .Lpatch1:
 ; CHECK-NEXT:  nop
 ;; Other basic blocks have BTI, but they don't affect our decision to not create .Lpatch0
 ; CHECK:      .LBB{{.+}} // %sw.bb1
 ; CHECK-NEXT:  hint #36
 ; CHECK:      .section __patchable_function_entries,"awo",@progbits,f1i{{$}}
 ; CHECK-NEXT: .p2align 3
-; CHECK-NEXT: .xword .Lfunc_begin3
+; CHECK-NEXT: .xword .Lpatch1
 entry:
   switch i64 %v, label %sw.bb0 [
     i64 1, label %sw.bb1