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

[MaskRay]

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.

[merge_guards_bot]

Unit tests: pass. 62315 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[MaskRay]

make ARCH=arm64 CROSS_COMPILE=aarch64-linux-gnu- CC=/tmp/ReleaseA/bin/clang LD=/tmp/ReleaseA/bin/ld.lld KCFLAGS=-mbranch-protection=bti O=/tmp/out/arm64 allmodconfig all -j 30

Picked one object file randomly. Verified that .rela__patchable_function_entries reference the instructions after BTI.

[nickdesaulniers]

Indeed, it looks like just from the emitted assembly GCC and LLVM will differ where they put the .section directives relative to the .text, but once assembled there's no difference from an object file perspective. I find putting the __patchable_function_entries outside of the function a nice touch. Thanks for the patch @MaskRay !

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

15

Would it be helpful have a test for patchable-function-entry > 1 and branch-target-enforcement?

[MaskRay]

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

15

I feel that may be redundant. Different numbers of NOPs are tested in patchable-function-entry.ll . For interaction with BTI, this one is sufficient.

[MaskRay]

@hans Please cherry pick 06b8e32d4fd3f634f793e3bc0bc4fdb885e7a3ac to release/10.x :)

[hans]

In D73680#1850170, @MaskRay wrote:

@hans Please cherry pick 06b8e32d4fd3f634f793e3bc0bc4fdb885e7a3ac to release/10.x :)

Done in cbd4815dec165fc4190f30cd6f409334c7e59063

llvm/include/llvm/CodeGen/AsmPrinter.h

public:
   MCSymbol *CurrentFnSymForSize = nullptr;
 
   /// Map global GOT equivalent MCSymbols to GlobalVariables and keep track of
   /// its number of uses by other globals.
   using GOTEquivUsePair = std::pair<const GlobalVariable *, unsigned>;
   MapVector<const MCSymbol *, GOTEquivUsePair> GlobalGOTEquivs;
 
 private:
-  MCSymbol *CurrentFnBegin = nullptr;
   MCSymbol *CurrentFnEnd = nullptr;
   MCSymbol *CurExceptionSym = nullptr;
 
   // The garbage collection metadata printer table.
   void *GCMetadataPrinters = nullptr; // Really a DenseMap.
 
   /// Emit comments in assembly output if this is true.
   bool VerboseAsm;
 
   static char ID;
 
 protected:
+  MCSymbol *CurrentFnBegin = nullptr;
+
   /// Protected struct HandlerInfo and Handlers permit target extended
   /// AsmPrinter adds their own handlers.
   struct HandlerInfo {
     std::unique_ptr<AsmPrinterHandler> Handler;
     const char *TimerName;
     const char *TimerDescription;
     const char *TimerGroupName;
     const char *TimerGroupDescription;

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

if (F.hasPrefixData()) {
     } else {
       EmitGlobalConstant(F.getParent()->getDataLayout(), F.getPrefixData());
     }
   }
 
   // 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;
   }
 
   // Emit the function descriptor. This is a virtual function to allow targets
   // to emit their specific function descriptor.
   if (MAI->needsFunctionDescriptors())
     EmitFunctionDescriptor();
 

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

void AArch64AsmPrinter::EmitInstruction(const MachineInstr *MI) {
   }
 
   AArch64TargetStreamer *TS =
     static_cast<AArch64TargetStreamer *>(OutStreamer->getTargetStreamer());
   // Do any manual lowerings.
   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);
       MachineOperand Hi_MOSym(MO_Sym), Lo_MOSym(MO_Sym);
       MCOperand Hi_MCSym, Lo_MCSym;
 
       Hi_MOSym.setTargetFlags(AArch64II::MO_G1 | AArch64II::MO_S);
       Lo_MOSym.setTargetFlags(AArch64II::MO_G0 | AArch64II::MO_NC);

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

 ; RUN: llc -mtriple=aarch64 %s -o - | FileCheck --check-prefixes=CHECK %s
 
 define void @f0() "patchable-function-entry"="0" "branch-target-enforcement" {
 ; CHECK-LABEL: f0:
 ; CHECK-NEXT: .Lfunc_begin0:
 ; CHECK:      // %bb.0:
 ; CHECK-NEXT:  hint #34
 ; CHECK-NEXT:  ret
 ; CHECK-NOT:  .section __patchable_function_entries
   ret void
 }
 
 ;; -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
 }
 
 ;; -fpatchable-function-entry=2,1 -mbranch-protection=bti
 define void @f2_1() "patchable-function-entry"="1" "patchable-function-prefix"="1" "branch-target-enforcement" {
 ; CHECK-LABEL: .type f2_1,@function
 ; CHECK-NEXT: .Ltmp0:
 ; CHECK-NEXT:  nop
 ; CHECK-NEXT: f2_1:
 ; CHECK-NEXT: .Lfunc_begin2:
 ; CHECK:      // %bb.0:
 ; CHECK-NEXT:  hint #34
 ; 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,f1,unique,0
 ; 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" {
+; 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
+}