[ARM,AArch64] Don't put BTI at asm goto branch targets (#141562)

In 'asm goto' statements ('callbr' in LLVM IR), you can specify one or
more labels / basic blocks in the containing function which the assembly
code might jump to. If you're also compiling with branch target
enforcement via BTI, then previously listing a basic block as a possible
jump destination of an asm goto would cause a BTI instruction to be
placed at the start of the block, in case the assembly code used an
_indirect_ branch instruction (i.e. to a destination address read from a
register) to jump to that location. Now it doesn't do that any more:
branches to destination labels from the assembly code are assumed to be
direct branches (to a relative offset encoded in the instruction), which
don't require a BTI at their destination.

This change was proposed in https://discourse.llvm.org/t/85845 and there
seemed to be no disagreement. The rationale is:

1. it brings clang's handling of asm goto in Arm and AArch64 in line
with gcc's, which didn't generate BTIs at the target labels in the first
place.

2. it improves performance in the Linux kernel, which uses a lot of 'asm
goto' in which the assembly language just contains a NOP, and the
label's address is saved elsewhere to let the kernel self-modify at run
time to swap between the original NOP and a direct branch to the label.
This allows hot code paths to be instrumented for debugging, at only the
cost of a NOP when the instrumentation is turned off, instead of the
larger cost of an indirect branch. In this situation a BTI is
unnecessary (if the branch happens it's direct), and since the code
paths are hot, also a noticeable performance hit.

Implementation:

`SelectionDAGBuilder::visitCallBr` is the place where 'asm goto' target
labels are handled. It calls `setIsInlineAsmBrIndirectTarget()` on each
target `MachineBasicBlock`. Previously it also called
`setMachineBlockAddressTaken()`, which made `hasAddressTaken()` return
true, which caused a BTI to be added in the Arm backends.

Now `visitCallBr` doesn't call `setMachineBlockAddressTaken()` any more
on asm goto targets, but `hasAddressTaken()` also checks the flag set by
`setIsInlineAsmBrIndirectTarget()`. So call sites that were using
`hasAddressTaken()` don't need to be modified. But the Arm backends
don't call `hasAddressTaken()` any more: instead they test two more
specific query functions that cover all the reasons `hasAddressTaken()`
might have returned true _except_ being an asm goto target.

Testing:

The new test `AArch64/callbr-asm-label-bti.ll` is testing the actual
change, where it expects not to see a `bti` instruction after
`[[LABEL]]`. The rest of the test changes are all churn, due to the
flags on basic blocks changing. Actual output code hasn't changed in any
of the existing tests, only comments and diagnostics.

Further work:

`RISCVIndirectBranchTracking.cpp` and `X86IndirectBranchTracking.cpp`
also call `hasAddressTaken()` in a way that might benefit from using the
same more specific check I've put in `ARMBranchTargets.cpp` and
`AArch64BranchTargets.cpp`. But I'm not sure of that, so in this commit
I've only changed the Arm backends, and left those alone.

Author: statham-arm

clang/docs/LanguageExtensions.rst

@@ -2043,6 +2043,17 @@ references can be used instead of numeric references.
       return -1;
   }
 
+ASM Goto versus Branch Target Enforcement
+=========================================
+
+Some target architectures implement branch target enforcement, by requiring
+indirect (register-controlled) branch instructions to jump only to locations
+marked by a special instruction (such as AArch64 ``bti``).
+
+The assembler code inside an ``asm goto`` statement is expected not to use a
+branch instruction of that kind to transfer control to any of its destination
+labels. Therefore, using a label in an ``asm goto`` statement does not cause
+clang to put a ``bti`` or equivalent instruction at the label.
 
 Constexpr strings in GNU ASM statements
 =======================================

clang/docs/ReleaseNotes.rst

@@ -41,6 +41,11 @@ Potentially Breaking Changes
 - For ARM targets when compiling assembly files, the features included in the selected CPU
   or Architecture's FPU are included. If you wish not to use a specific feature,
   the relevant ``+no`` option will need to be amended to the command line option.
+- When compiling with branch target enforcement, ``asm goto``
+  statements will no longer guarantee to place a ``bti`` or
+  ``endbr64`` instruction at the labels named as possible branch
+  destinations, so it is not safe to use a register-controlled branch
+  instruction to branch to one. (In line with gcc.)
 
 C/C++ Language Potentially Breaking Changes
 -------------------------------------------

llvm/docs/LangRef.rst

@@ -9596,6 +9596,14 @@ may not be equal to the address provided for the same block to this
 instruction's ``indirect labels`` operand. The assembly code may only transfer
 control to addresses provided via this instruction's ``indirect labels``.
 
+On target architectures that implement branch target enforcement by requiring
+indirect (register-controlled) branch instructions to jump only to locations
+marked by a special instruction (such as AArch64 ``bti``), the called code is
+expected not to use such an indirect branch to transfer control to the
+locations in ``indirect labels``. Therefore, including a label in the
+``indirect labels`` of a ``callbr`` does not require the compiler to put a
+``bti`` or equivalent instruction at the label.
+
 Arguments:
 """"""""""
 

llvm/docs/ReleaseNotes.md

@@ -70,6 +70,12 @@ Changes to the LLVM IR
   `llvm.load.relative`.
 * Inline asm calls no longer accept ``label`` arguments. Use ``callbr`` instead.
 
+* Updated semantics of the `callbr` instruction to clarify that its
+  'indirect labels' are not expected to be reached by indirect (as in
+  register-controlled) branch instructions, and therefore are not
+  guaranteed to start with a `bti` or `endbr64` instruction, where
+  those exist.
+
 Changes to LLVM infrastructure
 ------------------------------
 

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -277,18 +277,33 @@ class MachineBasicBlock
   /// of a terminator, exception-handling target, or jump table. This is
   /// either the result of an IR-level "blockaddress", or some form
   /// of target-specific branch lowering.
+  ///
+  /// The name of this function `hasAddressTaken` implies that the address of
+  /// the block is known and used in a general sense, but not necessarily that
+  /// the address is used by an indirect branch instruction. So branch target
+  /// enforcement need not put a BTI instruction (or equivalent) at the start
+  /// of a block just because this function returns true. The decision about
+  /// whether to add a BTI can be more subtle than that, and depends on the
+  /// more detailed checks that this function aggregates together.
   bool hasAddressTaken() const {
-    return MachineBlockAddressTaken || AddressTakenIRBlock;
+    return MachineBlockAddressTaken || AddressTakenIRBlock ||
+           IsInlineAsmBrIndirectTarget;
   }
 
   /// Test whether this block is used as something other than the target of a
   /// terminator, exception-handling target, jump table, or IR blockaddress.
   /// For example, its address might be loaded into a register, or
   /// stored in some branch table that isn't part of MachineJumpTableInfo.
+  ///
+  /// If this function returns true, it _does_ mean that branch target
+  /// enforcement needs to put a BTI or equivalent at the start of the block.
   bool isMachineBlockAddressTaken() const { return MachineBlockAddressTaken; }
 
   /// Test whether this block is the target of an IR BlockAddress.  (There can
   /// more than one MBB associated with an IR BB where the address is taken.)
+  ///
+  /// If this function returns true, it _does_ mean that branch target
+  /// enforcement needs to put a BTI or equivalent at the start of the block.
   bool isIRBlockAddressTaken() const { return AddressTakenIRBlock; }
 
   /// Retrieves the BasicBlock which corresponds to this MachineBasicBlock.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -4330,6 +4330,8 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
       OutStreamer->emitLabel(Sym);
   } else if (isVerbose() && MBB.isMachineBlockAddressTaken()) {
     OutStreamer->AddComment("Block address taken");
+  } else if (isVerbose() && MBB.isInlineAsmBrIndirectTarget()) {
+    OutStreamer->AddComment("Inline asm indirect target");
   }
 
   // Print some verbose block comments.

llvm/lib/CodeGen/BasicBlockPathCloning.cpp

@@ -121,14 +121,21 @@ bool IsValidCloning(const MachineFunction &MF,
       }
       if (PathBB->isMachineBlockAddressTaken()) {
         // Avoid cloning blocks which have their address taken since we can't
-        // rewire branches to those blocks as easily (e.g., branches within
-        // inline assembly).
+        // rewire branches to those blocks as easily.
         WithColor::warning()
             << "block #" << BBID
             << " has its machine block address taken in function "
             << MF.getName() << "\n";
         return false;
       }
+      if (PathBB->isInlineAsmBrIndirectTarget()) {
+        // Similarly for branches to the block within an asm goto.
+        WithColor::warning()
+            << "block #" << BBID
+            << " is a branch target of an 'asm goto' in function "
+            << MF.getName() << "\n";
+        return false;
+      }
     }
 
     if (I != ClonePath.size() - 1 && !PathBB->empty() &&

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -3399,7 +3399,11 @@ void SelectionDAGBuilder::visitCallBr(const CallBrInst &I) {
     BasicBlock *Dest = I.getIndirectDest(i);
     MachineBasicBlock *Target = FuncInfo.getMBB(Dest);
     Target->setIsInlineAsmBrIndirectTarget();
-    Target->setMachineBlockAddressTaken();
+    // If we introduce a type of asm goto statement that is permitted to use an
+    // indirect call instruction to jump to its labels, then we should add a
+    // call to Target->setMachineBlockAddressTaken() here, to mark the target
+    // block as requiring a BTI.
+
     Target->setLabelMustBeEmitted();
     // Don't add duplicate machine successors.
     if (Dests.insert(Dest).second)

llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

@@ -99,7 +99,8 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
 
     // If the block itself is address-taken, it could be indirectly branched
     // to, but not called.
-    if (MBB.hasAddressTaken() || JumpTableTargets.count(&MBB))
+    if (MBB.isMachineBlockAddressTaken() || MBB.isIRBlockAddressTaken() ||
+        JumpTableTargets.count(&MBB))
       CouldJump = true;
 
     if (CouldCall || CouldJump) {

llvm/lib/Target/ARM/ARMBranchTargets.cpp

@@ -77,7 +77,8 @@ bool ARMBranchTargets::runOnMachineFunction(MachineFunction &MF) {
     // modes. These modes do not support PACBTI. As a result, BTI instructions
     // are not added in the destination blocks.
 
-    if (IsFirstBB || MBB.hasAddressTaken() || MBB.isEHPad()) {
+    if (IsFirstBB || MBB.isMachineBlockAddressTaken() ||
+        MBB.isIRBlockAddressTaken() || MBB.isEHPad()) {
       addBTI(TII, MBB, IsFirstBB);
       MadeChange = true;
     }