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

[pbarrio]

Summary

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.

[pbarrio]

@MaskRay sorry to pull you into this review - feel free to ignore if you're not interested.

I'm removing one of your tests for patchable-function-entry, since my proposed changes mean all functions will start with BTI c. Just to double-check: is there any other scenario that you wanted to check with the removed test, other than a function with BTI j in the branche targets but no BTI c at the start of the function? Thanks!

[peter.smith]

This looks reasonable to me, while the linker moving a local function far away from its callers in the same object file is unlikely, it is at least theoretically possible with linker scripts or a symbol ordering file. My guess is that this would only affect a small number of functions?

llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

89

A linker is not permitted to add a thunk for the R_AARCH64_CONDBR19 relocation (A conditional branch), it may be possible to exclude the conditional branch in the check. My guess is that it won't be worth it as calling a function with a conditional branch is not likely to be generated by the compiler.

[MaskRay]

I know that for generalization, sometimes specific tests can be removed. -fpatchable-function-entry is a bit special though, it can place NOPs before and after the function label. This property is relied upon by the Linux kernel dynamic ftrace.
Since the NOP insertion is a bit strange, not clearly covered by other tests, and testing it does not increase test time (we don't add new llc commands so no process startup overhead), I'd like patchable-function-entry-bti.ll to be kept.

[pbarrio]

In D99417#2653048, @peter.smith wrote:

My guess is that this would only affect a small number of functions?

I got some numbers on the NDK samples (https://github.com/android/ndk-samples). Of a total of 6746 functions, this patch inserts a landing pad for 22 additional functions, which is ~0.3% of them. I haven't done proper benchmarking since this is a bugfix rather than an optimization, but I agree that this patch should affect a small number of functions.

[peter.smith]

In D99417#2658168, @pbarrio wrote:

In D99417#2653048, @peter.smith wrote:

My guess is that this would only affect a small number of functions?

I got some numbers on the NDK samples (https://github.com/android/ndk-samples). Of a total of 6746 functions, this patch inserts a landing pad for 22 additional functions, which is ~0.3% of them. I haven't done proper benchmarking since this is a bugfix rather than an optimization, but I agree that this patch should affect a small number of functions.

Thanks for the figures, that looks like we avoid a potential problem without much impact on the generated code.

For patchable functions are there any other cases where we wouldn't generate a BTI at the start of the function? If there is we may be alter the test so that it is still checking that the nop is still inserted. If there is no way to avoid a BTI then we may have to alter the expected result and the comment.

I've dug up https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92424 and https://reviews.llvm.org/D72215 for context.

[pbarrio]

Reintroduced the removed fpatchable-function-entry test, but modified so that it expects a landing pad at the beginning of the function.

After the changes, it seems @f1i tests the same thing as @f1, although I guess in a bigger function?

[pbarrio]

For patchable functions are there any other cases where we wouldn't generate a BTI at the start of the function? If there is we may be alter the test so that it is still checking that the nop is still inserted. If there is no way to avoid a BTI then we may have to alter the expected result and the comment.

I've dug up https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92424 and https://reviews.llvm.org/D72215 for context.

Thanks for the context! I don't think there is a way to generate a function without an initial landing pad ATM. The BTI placement pass will add a landing pad at the beginning of any BB that is the start of a function.

From a quick look at that GCC bug report, I think we should be fine now: the landing pads are place at the right point. If, at some point, the behaviour of LLVM changes again, I've reinstated the test that I initially removed, and hopefully it will shout.

[pbarrio]

Ping.

[peter.smith]

This LGTM, although please give a few days for MaskRay to check the patchable functions case.

[pbarrio]

Sure thing, thank you for the review. @MaskRay please let me know if you want to further discuss the changes to the patchable-function-entry test, otherwise I will commit this sometime next week.

[MaskRay]

LGTM.

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

48

The comment needs update. It should say that we add BTI c even if the function has the internal linkage.

[danielkiss]

LGTM

llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

 //===-- AArch64BranchTargets.cpp -- Harden code using v8.5-A BTI extension -==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 
 bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   if (!MF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
     return false;
 
   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
   // used for indirect branches, so need BTI instructions.
   SmallPtrSet<MachineBasicBlock *, 8> JumpTableTargets;
   if (auto *JTI = MF.getJumpTableInfo())
     for (auto &JTE : JTI->getJumpTables())
       for (auto *MBB : JTE.MBBs)
         JumpTableTargets.insert(MBB);
 
   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
     // to, but not called.
     if (MBB.hasAddressTaken() || JumpTableTargets.count(&MBB))
       CouldJump = true;
 
     if (CouldCall || CouldJump) {

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

body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: simple_external
     ; CHECK: HINT 34
     ; CHECK: RET
     $w0 = ORRWrs $wzr, $wzr, 0
     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
 
 ---
 # Function starts with PACIASP, which implicitly acts as BTI JC, so no change
 # needed.
 name:            ptr_auth

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

 ; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ; CHECK:      .Lfunc_end2:
 ; CHECK-NEXT: .size f2_1, .Lfunc_end2-f2_1
 ; CHECK:      .section __patchable_function_entries,"awo",@progbits,f2_1{{$}}
 ; CHECK-NEXT: .p2align 3
 ; CHECK-NEXT: .xword .Ltmp0
   ret void
 }
-
-;; -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.
-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:  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
-entry:
-  switch i64 %v, label %sw.bb0 [
-    i64 1, label %sw.bb1
-    i64 2, label %sw.bb2
-    i64 3, label %sw.bb3
-    i64 4, label %sw.bb4
-  ]
-sw.bb0:
-  call void asm sideeffect "", ""()
-  ret void
-sw.bb1:
-  call void asm sideeffect "", ""()
-  ret void
-sw.bb2:
-  call void asm sideeffect "", ""()
-  ret void
-sw.bb3:
-  call void asm sideeffect "", ""()
-  ret void
-sw.bb4:
-  call void asm sideeffect "", ""()
-  ret void
-}