[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

2752

"No newline at end of file"

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 target-independent 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;
+          }
+          emitNops(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++;

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/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()) {
     // after it.
     SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
     // Then emit the call
     OutStreamer->emitInstruction(TmpInst, getSubtargetInfo());
     return;
   }
 
   EmitAndCountInstruction(TmpInst);
 }

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
+}
 No newline at end of file

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

 
 ;; Without -function-sections, f2 is in the same text section as f1.
 ;; They share the __patchable_function_entries section.
 ;; With -function-sections, f1 and f2 are in different text sections.
 ;; Use separate __patchable_function_entries.
 define void @f2() "patchable-function-entry"="2" {
 ; CHECK-LABEL: f2:
 ; CHECK-NEXT: .Lfunc_begin2:
-; 32:          xchgw %ax, %ax
-; 64:          xchgw %ax, %ax
+; CHECK:  nop
+; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ; CHECK:       .section __patchable_function_entries,"awo",@progbits,f2{{$}}
 ; 32:          .p2align 2
 ; 32-NEXT:     .long .Lfunc_begin2
 ; 64:          .p2align 3
 ; 64-NEXT:     .quad .Lfunc_begin2
   ret void
 }
 
 $f3 = comdat any
 define void @f3() "patchable-function-entry"="3" comdat {
 ; CHECK-LABEL: f3:
 ; CHECK-NEXT: .Lfunc_begin3:
-; 32:          xchgw %ax, %ax
-; 32-NEXT:     nop
-; 64:          nopl (%rax)
+; CHECK:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
 ; CHECK:       ret
 ; CHECK:       .section __patchable_function_entries,"aGwo",@progbits,f3,comdat,f3{{$}}
 ; 32:          .p2align 2
 ; 32-NEXT:     .long .Lfunc_begin3
 ; 64:          .p2align 3
 ; 64-NEXT:     .quad .Lfunc_begin3
   ret void
 }
 
 $f5 = comdat any
 define void @f5() "patchable-function-entry"="5" comdat {
 ; CHECK-LABEL: f5:
 ; CHECK-NEXT: .Lfunc_begin4:
-; 32-COUNT-2:  xchgw %ax, %ax
-; 32-NEXT:     nop
-; 64:          nopl 8(%rax,%rax)
+; CHECK:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ; CHECK:       .section __patchable_function_entries,"aGwo",@progbits,f5,comdat,f5{{$}}
 ; 32:          .p2align 2
 ; 32-NEXT:     .long .Lfunc_begin4
 ; 64:          .p2align 3
 ; 64-NEXT:     .quad .Lfunc_begin4
   ret void
 }
 
 ;; -fpatchable-function-entry=3,2
 ;; "patchable-function-prefix" emits data before the function entry label.
 ;; We emit 1-byte NOPs before the function entry, so that with a partial patch,
 ;; the remaining instructions do not need to be modified.
 define void @f3_2() "patchable-function-entry"="1" "patchable-function-prefix"="2" {
 ; CHECK-LABEL: .type f3_2,@function
 ; CHECK-NEXT: .Ltmp0: # @f3_2
-; CHECK-NEXT:  nop
+; CHECK:  nop
 ; CHECK-NEXT:  nop
 ; CHECK-NEXT: f3_2:
 ; CHECK:      # %bb.0:
 ; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ;; .size does not include the prefix.
 ; CHECK:      .Lfunc_end5:
 ; CHECK-NEXT: .size f3_2, .Lfunc_end5-f3_2
 ; CHECK:      .section __patchable_function_entries,"awo",@progbits,f3_2{{$}}
 ; 32:         .p2align 2
 ; 32-NEXT:    .long .Ltmp0
 ; 64:         .p2align 3
 ; 64-NEXT:    .quad .Ltmp0
   %frame = alloca i8, i32 16
   ret void
 }