diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h --- a/llvm/include/llvm/CodeGen/AsmPrinter.h +++ b/llvm/include/llvm/CodeGen/AsmPrinter.h @@ -135,7 +135,6 @@ MapVector GlobalGOTEquivs; private: - MCSymbol *CurrentFnBegin = nullptr; MCSymbol *CurrentFnEnd = nullptr; MCSymbol *CurExceptionSym = nullptr; @@ -148,6 +147,8 @@ 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 --- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp +++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp @@ -727,15 +727,21 @@ // 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 --- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -1000,6 +1000,26 @@ 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 --- a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll +++ b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll @@ -11,15 +11,18 @@ } ;; -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 @@ ; 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 +}