[AArch64] Add function attribute "patchable-function-entry" to add NOPs at function entry

The Linux kernel uses -fpatchable-function-entry to implement DYNAMIC_FTRACE_WITH_REGS
for arm64 and parisc. GCC 8 implemented
-fpatchable-function-entry, which can be seen as a generalized form of
-mnop-mcount. The N,M form (function entry points before the Mth NOP) is
currently only used by parisc.

This patch adds N,0 support to AArch64 codegen. N is represented as the
function attribute "patchable-function-entry". We will use a different
function attribute for M, if we decide to implement it.

The patch reuses the existing patchable-function pass, and
TargetOpcode::PATCHABLE_FUNCTION_ENTER which is currently used by XRay.

When the integrated assembler is used, __patchable_function_entries will
be created for each text section with the SHF_LINK_ORDER flag to prevent
--gc-sections (https://gcc.gnu.org/bugzilla/show_bug.cgi?id=93197) and
COMDAT (https://gcc.gnu.org/bugzilla/show_bug.cgi?id=93195) issues.

Retrospectively, __patchable_function_entries should use a PC-relative
relocation type to avoid the SHF_WRITE flag and dynamic relocations.

"patchable-function-entry"'s interaction with Branch Target
Identification is still unclear (see
https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92424 for GCC discussions).

Reviewed By: peter.smith

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

4 files changed, 70 insertions, 2 deletions

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 9f68b07b76c..1579646352d 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1222,6 +1222,8 @@ void AsmPrinter::EmitFunctionBody() {
   // Emit section containing stack size metadata.
   emitStackSizeSection(*MF);
 
+  emitPatchableFunctionEntries();
+
   if (isVerbose())
     OutStreamer->GetCommentOS() << "-- End function\n";
 
@@ -1660,6 +1662,7 @@ MCSymbol *AsmPrinter::getCurExceptionSym() {
 
 void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
   this->MF = &MF;
+  const Function &F = MF.getFunction();
 
   // Get the function symbol.
   if (MAI->needsFunctionDescriptors()) {
@@ -1672,7 +1675,6 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
     CurrentFnSym =
         OutContext.getOrCreateSymbol("." + CurrentFnDescSym->getName());
 
-    const Function &F = MF.getFunction();
     MCSectionXCOFF *FnEntryPointSec =
         cast<MCSectionXCOFF>(getObjFileLowering().SectionForGlobal(&F, TM));
     // Set the containing csect.
@@ -1685,7 +1687,8 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
   CurrentFnBegin = nullptr;
   CurExceptionSym = nullptr;
   bool NeedsLocalForSize = MAI->needsLocalForSize();
-  if (needFuncLabelsForEHOrDebugInfo(MF, MMI) || NeedsLocalForSize ||
+  if (F.hasFnAttribute("patchable-function-entry") ||
+      needFuncLabelsForEHOrDebugInfo(MF, MMI) || NeedsLocalForSize ||
       MF.getTarget().Options.EmitStackSizeSection) {
     CurrentFnBegin = createTempSymbol("func_begin");
     if (NeedsLocalForSize)
@@ -3194,6 +3197,38 @@ void AsmPrinter::recordSled(MCSymbol *Sled, const MachineInstr &MI,
                                        AlwaysInstrument, &F, Version});
 }
 
+void AsmPrinter::emitPatchableFunctionEntries() {
+  const Function &F = MF->getFunction();
+  if (!F.hasFnAttribute("patchable-function-entry"))
+    return;
+  const unsigned PointerSize = getPointerSize();
+  if (TM.getTargetTriple().isOSBinFormatELF()) {
+    auto Flags = ELF::SHF_WRITE | ELF::SHF_ALLOC;
+    std::string GroupName;
+    if (F.hasComdat()) {
+      Flags |= ELF::SHF_GROUP;
+      GroupName = F.getComdat()->getName();
+    }
+
+    // As of binutils 2.33, GNU as does not support section flag "o". Use
+    // SHF_LINK_ORDER if we are using the integrated assembler.
+    MCSymbolELF *Link = MAI->useIntegratedAssembler()
+                            ? cast<MCSymbolELF>(CurrentFnSym)
+                            : nullptr;
+    if (Link)
+      Flags |= ELF::SHF_LINK_ORDER;
+
+    MCSection *Section = getObjFileLowering().SectionForGlobal(&F, TM);
+    auto R = PatchableFunctionEntryID.try_emplace(
+        Section, PatchableFunctionEntryID.size());
+    OutStreamer->SwitchSection(OutContext.getELFSection(
+        "__patchable_function_entries", ELF::SHT_PROGBITS, Flags, 0, GroupName,
+        R.first->second, Link));
+    EmitAlignment(Align(PointerSize));
+    OutStreamer->EmitSymbolValue(CurrentFnBegin, PointerSize);
+  }
+}
+
 uint16_t AsmPrinter::getDwarfVersion() const {
   return OutStreamer->getContext().getDwarfVersion();
 }
diff --git a/llvm/lib/CodeGen/PatchableFunction.cpp b/llvm/lib/CodeGen/PatchableFunction.cpp
index f499241d43c..1d6069c5055 100644
--- a/llvm/lib/CodeGen/PatchableFunction.cpp
+++ b/llvm/lib/CodeGen/PatchableFunction.cpp
@@ -55,6 +55,15 @@ static bool doesNotGeneratecode(const MachineInstr &MI) {
 }
 
 bool PatchableFunction::runOnMachineFunction(MachineFunction &MF) {
+  if (MF.getFunction().hasFnAttribute("patchable-function-entry")) {
+    MachineBasicBlock &FirstMBB = *MF.begin();
+    MachineInstr &FirstMI = *FirstMBB.begin();
+    const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+    BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
+            TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
+    return true;
+  }
+
   if (!MF.getFunction().hasFnAttribute("patchable-function"))
     return false;
 
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index d232946af29..d15b70d71b4 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -1851,6 +1851,18 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
     if (FP != "all" && FP != "non-leaf" && FP != "none")
       CheckFailed("invalid value for 'frame-pointer' attribute: " + FP, V);
   }
+
+  if (Attrs.hasFnAttribute("patchable-function-entry")) {
+    StringRef S0 = Attrs
+                       .getAttribute(AttributeList::FunctionIndex,
+                                     "patchable-function-entry")
+                       .getValueAsString();
+    StringRef S = S0;
+    unsigned N;
+    if (S.getAsInteger(10, N))
+      CheckFailed(
+          "\"patchable-function-entry\" takes an unsigned integer: " + S0, V);
+  }
 }
 
 void Verifier::verifyFunctionMetadata(
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 2026bcce3e8..e062ca77e71 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -243,6 +243,18 @@ void AArch64AsmPrinter::EmitStartOfAsmFile(Module &M) {
 
 void AArch64AsmPrinter::LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI)
 {
+  const Function &F = MF->getFunction();
+  if (F.hasFnAttribute("patchable-function-entry")) {
+    unsigned Num;
+    if (F.getFnAttribute("patchable-function-entry")
+            .getValueAsString()
+            .getAsInteger(10, Num))
+      return;
+    for (; Num; --Num)
+      EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::HINT).addImm(0));
+    return;
+  }
+
   EmitSled(MI, SledKind::FUNCTION_ENTER);
 }