[XRay] Add generic patchable-function-entry NOP sled implementation

[duck-37]

Summary

This patch generalizes the logic for patchable-function-entry=n NOP sleds, meaning that this attribute now has consistent behavior across all architectures that support XRay. The differences in codegen was noticed while porting a separate test-case over to ARM32.

Closes https://github.com/llvm/llvm-project/issues/60672

[MaskRay]

The increased instruction count for x86 is a slight performance regression. Can it be fixed?

[MaskRay]

Fixes https://github.com/llvm/llvm-project/issues/60672

It's not a bug so I'd use Close. Note: without runtime support, the compiler codegen change isn't really useful

[duck-37]

In D143802#4119979, @MaskRay wrote:

Fixes https://github.com/llvm/llvm-project/issues/60672

It's not a bug so I'd use Close. Note: without runtime support, the compiler codegen change isn't really useful

Sure, I'll adjust that for the next revision. The main idea here was consistency across architectures; I came across this while moving a test-case over to ARM32 and noticing a difference in the behavior.

[duck-37]

Changed some wording in the review, used a virtual function instead of hardcoded emitNops call so that X86 (and potentially other targets) can use their own implementations if necessary.

[MaskRay]

Sorry for the delay. This looks good, but it needs rebase:)

llvm/lib/Target/X86/X86MCInstLower.cpp

2756

"No newline at end of file"

llvm/include/llvm/CodeGen/AsmPrinter.h

public:
   void recordSled(MCSymbol *Sled, const MachineInstr &MI, SledKind Kind,
                   uint8_t Version = 0);
 
   /// Emit a table with all XRay instrumentation points.
   void emitXRayTable();
 
   void emitPatchableFunctionEntries();
 
+  /// Emit a simple NOP sled.
+  virtual void emitPatchableNOPSled(unsigned Num);
+
   //===------------------------------------------------------------------===//
   // MachineFunctionPass Implementation.
   //===------------------------------------------------------------------===//
 
   /// Record analysis usage.
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   /// Set up the AsmPrinter when we are working on a new module. If your pass

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

for (auto &MI : MBB) {
         break;
       case TargetOpcode::ARITH_FENCE:
         if (isVerbose())
           OutStreamer->emitRawComment("ARITH_FENCE");
         break;
       case TargetOpcode::MEMBARRIER:
         OutStreamer->emitRawComment("MEMBARRIER");
         break;
+
+      // patchable-function-entry=N is a simple NOP sled.
+      case TargetOpcode::PATCHABLE_FUNCTION_ENTER: {
+        const Function &F = MF->getFunction();
+        if (F.hasFnAttribute("patchable-function-entry")) {
+          unsigned Num;
+          if (F.getFnAttribute("patchable-function-entry")
+                  .getValueAsString()
+                  .getAsInteger(10, Num)) {
+            // This is garbage, do nothing.
+            break;
+          }
+          emitPatchableNOPSled(Num);
+        } else {
+          emitInstruction(&MI);
+        }
+        break;
+      }
+
       default:
         emitInstruction(&MI);
         if (CanDoExtraAnalysis) {
           MCInst MCI;
           MCI.setOpcode(MI.getOpcode());
           auto Name = OutStreamer->getMnemonic(MCI);
           auto I = MnemonicCounts.insert({Name, 0u});
           I.first->second++;
if (TM.getTargetTriple().isOSBinFormatELF()) {
     OutStreamer->switchSection(OutContext.getELFSection(
         "__patchable_function_entries", ELF::SHT_PROGBITS, Flags, 0, GroupName,
         F.hasComdat(), MCSection::NonUniqueID, LinkedToSym));
     emitAlignment(Align(PointerSize));
     OutStreamer->emitSymbolValue(CurrentPatchableFunctionEntrySym, PointerSize);
   }
 }
 
+void AsmPrinter::emitPatchableNOPSled(unsigned Num) {
+  // For most architectures, we just emit a set number of NOPs.
+  emitNops(Num);
+}
+
 uint16_t AsmPrinter::getDwarfVersion() const {
   return OutStreamer->getContext().getDwarfVersion();
 }
 
 void AsmPrinter::setDwarfVersion(uint16_t Version) {
   OutStreamer->getContext().setDwarfVersion(Version);
 }
 

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

void AArch64AsmPrinter::emitFunctionHeaderComment() {
   const AArch64FunctionInfo *FI = MF->getInfo<AArch64FunctionInfo>();
   std::optional<std::string> OutlinerString = FI->getOutliningStyle();
   if (OutlinerString != std::nullopt)
     OutStreamer->getCommentOS() << ' ' << OutlinerString;
 }
 
 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;
-    emitNops(Num);
-    return;
-  }
-
   emitSled(MI, SledKind::FUNCTION_ENTER);
 }
 
 void AArch64AsmPrinter::LowerPATCHABLE_FUNCTION_EXIT(const MachineInstr &MI) {
   emitSled(MI, SledKind::FUNCTION_EXIT);
 }
 
 void AArch64AsmPrinter::LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI) {

llvm/lib/Target/RISCV/RISCVMCInstLower.cpp

bool llvm::lowerRISCVMachineInstrToMCInst(const MachineInstr *MI, MCInst &OutMI,
 
   for (const MachineOperand &MO : MI->operands()) {
     MCOperand MCOp;
     if (lowerRISCVMachineOperandToMCOperand(MO, MCOp, AP))
       OutMI.addOperand(MCOp);
   }
 
   switch (OutMI.getOpcode()) {
-  case TargetOpcode::PATCHABLE_FUNCTION_ENTER: {
-    const Function &F = MI->getParent()->getParent()->getFunction();
-    if (F.hasFnAttribute("patchable-function-entry")) {
-      unsigned Num;
-      if (F.getFnAttribute("patchable-function-entry")
-              .getValueAsString()
-              .getAsInteger(10, Num))
-        return false;
-      AP.emitNops(Num);
-      return true;
-    }
-    break;
-  }
   case RISCV::PseudoReadVLENB:
     OutMI.setOpcode(RISCV::CSRRS);
     OutMI.addOperand(MCOperand::createImm(
         RISCVSysReg::lookupSysRegByName("VLENB")->Encoding));
     OutMI.addOperand(MCOperand::createReg(RISCV::X0));
     break;
   case RISCV::PseudoReadVL:
     OutMI.setOpcode(RISCV::CSRRS);
     OutMI.addOperand(
         MCOperand::createImm(RISCVSysReg::lookupSysRegByName("VL")->Encoding));
     OutMI.addOperand(MCOperand::createReg(RISCV::X0));
     break;
   }
   return false;
 }

llvm/lib/Target/X86/X86AsmPrinter.h

public:
   bool runOnMachineFunction(MachineFunction &MF) override;
   void emitFunctionBodyStart() override;
   void emitFunctionBodyEnd() override;
   void emitKCFITypeId(const MachineFunction &MF) override;
 
   bool shouldEmitWeakSwiftAsyncExtendedFramePointerFlags() const override {
     return ShouldEmitWeakSwiftAsyncExtendedFramePointerFlags;
   }
+
+  void emitPatchableNOPSled(unsigned Num) override;
 };
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/X86/X86MCInstLower.cpp

void X86AsmPrinter::LowerPATCHABLE_TYPED_EVENT_CALL(const MachineInstr &MI,
   recordSled(CurSled, MI, SledKind::TYPED_EVENT, 2);
 }
 
 void X86AsmPrinter::LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI,
                                                   X86MCInstLower &MCIL) {
 
   NoAutoPaddingScope NoPadScope(*OutStreamer);
 
-  const Function &F = MF->getFunction();
-  if (F.hasFnAttribute("patchable-function-entry")) {
-    unsigned Num;
-    if (F.getFnAttribute("patchable-function-entry")
-            .getValueAsString()
-            .getAsInteger(10, Num))
-      return;
-    emitX86Nops(*OutStreamer, Num, Subtarget);
-    return;
-  }
   // We want to emit the following pattern:
   //
   //   .p2align 1, ...
   // .Lxray_sled_N:
   //   jmp .tmpN
   //   # 9 bytes worth of noops
   //
   // We need the 9 bytes because at runtime, we'd be patching over the full 11
if (MI->isCall()) {
     SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
     // Then emit the call
     OutStreamer->emitInstruction(TmpInst, getSubtargetInfo());
     return;
   }
 
   EmitAndCountInstruction(TmpInst);
 }
+
+void X86AsmPrinter::emitPatchableNOPSled(unsigned Num) {
+  emitX86Nops(*OutStreamer, Num, Subtarget);
+}
 No newline at end of file

llvm/test/CodeGen/ARM/patchable-function-entry.ll

+; RUN: llc -mtriple=arm-linux-gnueabi -verify-machineinstrs < %s | FileCheck %s
+
+define void @f0() "patchable-function-entry"="0" {
+; CHECK-LABEL: f0:
+; CHECK-NOT: mov     r0, r0
+; CHECK:     mov     pc, lr
+  ret void
+}
+
+define void @f1() "patchable-function-entry"="1" {
+; CHECK-LABEL: f1:
+; CHECK:     mov     r0, r0
+; CHECK-NOT: mov     r0, r0
+; CHECK:     mov     pc, lr
+  ret void
+}
+
+define void @f2() "patchable-function-entry"="2" {
+; CHECK-LABEL: f2:
+; CHECK:     mov     r0, r0
+; CHECK:     mov     r0, r0
+; CHECK-NOT: mov     r0, r0
+; CHECK:     mov     pc, lr
+  ret void
+}

llvm/test/CodeGen/Hexagon/patchable-function-entry.ll

+; RUN: llc -march=hexagon -verify-machineinstrs < %s | FileCheck %s
+
+define void @f0() "patchable-function-entry"="0" {
+; CHECK-LABEL: f0:
+; CHECK-NOT: nop
+; CHECK: jumpr r31
+  ret void
+}
+
+define void @f1() "patchable-function-entry"="1" {
+; CHECK-LABEL: f1:
+; CHECK: nop
+; CHECK-NOT: nop
+; CHECK: jumpr r31
+  ret void
+}
+
+define void @f2() "patchable-function-entry"="2" {
+; CHECK-LABEL: f2:
+; CHECK: nop
+; CHECK: nop
+; CHECK-NOT: nop
+; CHECK: jumpr r31
+  ret void
+}