[AArch64] -fpatchable-function-entry=N,0: place patch label after BTI

Summary:
For -fpatchable-function-entry=N,0 -mbranch-protection=bti, after
9a24488cb67a90f889529987275c5e411ce01dda, we place the NOP sled after
the initial BTI.

```
.Lfunc_begin0:
bti c
nop
nop

.section __patchable_function_entries,"awo",@progbits,f,unique,0
.p2align 3
.xword .Lfunc_begin0
```

This patch adds a label after the initial BTI and changes the __patchable_function_entries entry to reference the label:

```
.Lfunc_begin0:
bti c
.Lpatch0:
nop
nop

.section __patchable_function_entries,"awo",@progbits,f,unique,0
.p2align 3
.xword .Lpatch0
```

This placement is compatible with the resolution in
https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92424 .

A local linkage function whose address is not taken does not need a BTI.
Placing the patch label after BTI has the advantage that code does not
need to differentiate whether the function has an initial BTI.

Reviewers: mrutland, nickdesaulniers, nsz, ostannard

Subscribers: kristof.beyls, hiraditya, llvm-commits

Tags: #llvm

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

(cherry picked from commit 06b8e32d4fd3f634f793e3bc0bc4fdb885e7a3ac)

4 files changed, 72 insertions, 4 deletions

diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index a860ce2773e..c710c5d7055 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -135,7 +135,6 @@ public:
   MapVector<const MCSymbol *, GOTEquivUsePair> GlobalGOTEquivs;
 
 private:
-  MCSymbol *CurrentFnBegin = nullptr;
   MCSymbol *CurrentFnEnd = nullptr;
   MCSymbol *CurExceptionSym = nullptr;
 
@@ -148,6 +147,8 @@ private:
   static char ID;
 
 protected:
+  MCSymbol *CurrentFnBegin = nullptr;
+
   /// Protected struct HandlerInfo and Handlers permit target extended
   /// AsmPrinter adds their own handlers.
   struct HandlerInfo {
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 3516f4a7b37..20cd9da31fb 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -709,15 +709,21 @@ void AsmPrinter::EmitFunctionHeader() {
   // Emit M NOPs for -fpatchable-function-entry=N,M where M>0. We arbitrarily
   // place prefix data before NOPs.
   unsigned PatchableFunctionPrefix = 0;
+  unsigned PatchableFunctionEntry = 0;
   (void)F.getFnAttribute("patchable-function-prefix")
       .getValueAsString()
       .getAsInteger(10, PatchableFunctionPrefix);
+  (void)F.getFnAttribute("patchable-function-entry")
+      .getValueAsString()
+      .getAsInteger(10, PatchableFunctionEntry);
   if (PatchableFunctionPrefix) {
     CurrentPatchableFunctionEntrySym =
         OutContext.createLinkerPrivateTempSymbol();
     OutStreamer->EmitLabel(CurrentPatchableFunctionEntrySym);
     emitNops(PatchableFunctionPrefix);
-  } else {
+  } else if (PatchableFunctionEntry) {
+    // May be reassigned when emitting the body, to reference the label after
+    // the initial BTI (AArch64) or endbr32/endbr64 (x86).
     CurrentPatchableFunctionEntrySym = CurrentFnBegin;
   }
 
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index b8953583a31..6da089d1859 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -1000,6 +1000,26 @@ void AArch64AsmPrinter::EmitInstruction(const MachineInstr *MI) {
   switch (MI->getOpcode()) {
   default:
     break;
+  case AArch64::HINT: {
+    // CurrentPatchableFunctionEntrySym can be CurrentFnBegin only for
+    // -fpatchable-function-entry=N,0. The entry MBB is guaranteed to be
+    // non-empty. If MI is the initial BTI, place the
+    // __patchable_function_entries label after BTI.
+    if (CurrentPatchableFunctionEntrySym &&
+        CurrentPatchableFunctionEntrySym == CurrentFnBegin &&
+        MI == &MF->front().front()) {
+      int64_t Imm = MI->getOperand(0).getImm();
+      if ((Imm & 32) && (Imm & 6)) {
+        MCInst Inst;
+        MCInstLowering.Lower(MI, Inst);
+        EmitToStreamer(*OutStreamer, Inst);
+        CurrentPatchableFunctionEntrySym = createTempSymbol("patch");
+        OutStreamer->EmitLabel(CurrentPatchableFunctionEntrySym);
+        return;
+      }
+    }
+    break;
+  }
     case AArch64::MOVMCSym: {
       Register DestReg = MI->getOperand(0).getReg();
       const MachineOperand &MO_Sym = MI->getOperand(1);
diff --git a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
index 8d889a73f9f..4b85664673f 100644
--- a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
+++ b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
@@ -11,15 +11,18 @@ define void @f0() "patchable-function-entry"="0" "branch-target-enforcement" {
 }
 
 ;; -fpatchable-function-entry=1 -mbranch-protection=bti
+;; For M=0, place the label .Lpatch0 after the initial BTI.
 define void @f1() "patchable-function-entry"="1" "branch-target-enforcement" {
 ; CHECK-LABEL: f1:
 ; CHECK-NEXT: .Lfunc_begin1:
-; CHECK:       hint #34
+; CHECK:      // %bb.0:
+; CHECK-NEXT:  hint #34
+; CHECK-NEXT: .Lpatch0:
 ; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ; CHECK:      .section __patchable_function_entries,"awo",@progbits,f1,unique,0
 ; CHECK-NEXT: .p2align 3
-; CHECK-NEXT: .xword .Lfunc_begin1
+; CHECK-NEXT: .xword .Lpatch0
   ret void
 }
 
@@ -41,3 +44,41 @@ define void @f2_1() "patchable-function-entry"="1" "patchable-function-prefix"="
 ; 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" {
+; 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,f1,unique,0
+; 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
+}