[AArch64] [XRay] Account for XRay event instrs in Branch Relaxation

[dhoekwater]

Summary

PATCHABLE_TYPED_EVENT_CALL and PATCHABLE_EVENT_CALL are pseudo
instructions that expand to XRay sleds, so getInstSizeInBytes
should reflect the size of the sleds, not the pseudo-instructions.

[dhoekwater]

Add a test

[paquette]

Can you add comments explaining the size in bytes?

"This expands to an XRay sled can be at most N instructions = M bytes"

Or something

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

131

comment?

143

comment?

[dhoekwater]

Add comments explaining the size of XRay PATCHABLE instructions

[dhoekwater]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

131

Done! Thanks for the feedback.

143

Done!

[dhoekwater]

Better spacing with clang-format

[paquette]

LGTM

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

: AArch64GenInstrInfo(AArch64::ADJCALLSTACKDOWN, AArch64::ADJCALLSTACKUP,
                           AArch64::CATCHRET),
       RI(STI.getTargetTriple()), Subtarget(STI) {}
 
 /// GetInstSize - Return the number of bytes of code the specified
 /// instruction may be.  This returns the maximum number of bytes.
 unsigned AArch64InstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   const MachineBasicBlock &MBB = *MI.getParent();
   const MachineFunction *MF = MBB.getParent();
+  const Function &F = MF->getFunction();
   const MCAsmInfo *MAI = MF->getTarget().getMCAsmInfo();
 
   {
     auto Op = MI.getOpcode();
     if (Op == AArch64::INLINEASM || Op == AArch64::INLINEASM_BR)
       return getInlineAsmLength(MI.getOperand(0).getSymbolName(), *MAI);
   }
 
unsigned AArch64InstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   case TargetOpcode::STATEPOINT:
     NumBytes = StatepointOpers(&MI).getNumPatchBytes();
     assert(NumBytes % 4 == 0 && "Invalid number of NOP bytes requested!");
     // No patch bytes means a normal call inst is emitted
     if (NumBytes == 0)
       NumBytes = 4;
     break;
   case TargetOpcode::PATCHABLE_FUNCTION_ENTER:
+    // If `patchable-function-entry` is set, PATCHABLE_FUNCTION_ENTER
+    // instructions are expanded to the specified number of NOPs. Otherwise,
+    // they are expanded to 36-byte XRay sleds.
+    NumBytes =
+        F.getFnAttributeAsParsedInteger("patchable-function-entry", 9) * 4;
+    break;
   case TargetOpcode::PATCHABLE_FUNCTION_EXIT:
+  case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
     // An XRay sled can be 4 bytes of alignment plus a 32-byte block.
     NumBytes = 36;
     break;
+  case TargetOpcode::PATCHABLE_EVENT_CALL:
+    // EVENT_CALL XRay sleds are exactly 6 instructions long (no alignment).
+    NumBytes = 24;
+    break;
 
   case AArch64::SPACE:
     NumBytes = MI.getOperand(1).getImm();
     break;
   case TargetOpcode::BUNDLE:
     NumBytes = getInstBundleLength(MI);
     break;
   }

llvm/test/CodeGen/AArch64/branch-relax-xray.ll

-; RUN: llc -mtriple=aarch64-unknown-linux-gnu -aarch64-tbz-offset-bits=4 -aarch64-cbz-offset-bits=4 < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -aarch64-tbz-offset-bits=4 -aarch64-cbz-offset-bits=3 < %s | FileCheck %s
 
 ;; Check that branch relaxation accounts for the size of xray EXIT sleds
 ;; Note that TAIL_CALL sleds don't exist on AArch64 and don't need a test.
 define void @exit(i1 zeroext %0) nounwind "function-instrument"="xray-always" {
 ; CHECK-LABEL: exit:
 ; CHECK-NEXT:     .Lfunc_begin0:
 ; CHECK-NEXT:     // %bb.0:
 ; CHECK-NEXT:       .p2align 2
end1:
   %2 = call i32 @bar()
   ret void
 
 end2:
   %3 = call i32 @baz()
   ret void
 }
 
+;; Check that branch relaxation accounts for the size of xray EVENT sleds
+define void @customevent(i1 zeroext %0, ptr nocapture noundef readonly %e1, i64 noundef %s1, ptr nocapture noundef readonly %e2, i64 noundef %s2) "function-instrument"="xray-always" {
+; CHECK-LABEL: customevent:
+; CHECK-NEXT:     .Lfunc_begin1:
+; CHECK-NEXT:       .cfi_startproc
+; CHECK-NEXT:       // %bb.0:
+; CHECK-NEXT:       .p2align 2
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK:            cbnz
+; CHECK-SAME:           [[FALLTHROUGH_2:.LBB[0-9_]+]]
+; CHECK-NEXT:       b
+; CHECK-SAME:           [[OUT_OF_RANGE_2:.LBB[0-9_]+]]
+; CHECK-NEXT:     [[FALLTHROUGH_2]]:
+; CHECK-SAME:                        // %end1
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK-NEXT:         Begin XRay custom event
+; CHECK:            bl      __xray_CustomEvent
+; CHECK:              End XRay custom event
+; CHECK-NEXT:     [[OUT_OF_RANGE_2]]:
+; CHECK-SAME:                        // %end2
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK-NEXT:         Begin XRay custom event
+; CHECK:            bl      __xray_CustomEvent
+; CHECK:              End XRay custom event
+; CHECK:          .Ltmp
+; CHECK-NEXT:       ret
+entry:
+  br i1 %0, label %end1, label %end2
+
+end1:
+  call void @llvm.xray.customevent(ptr %e1, i64 %s1)
+  br label %end2
+
+end2:
+  tail call void @llvm.xray.customevent(ptr %e2, i64 %s2)
+  ret void
+}
+
+;; Check that branch relaxation accounts for the size of xray TYPED_EVENT sleds
+define void @typedevent(i1 zeroext %0, i64 noundef %type, ptr nocapture noundef readonly %event, i64 noundef %size) "function-instrument"="xray-always" {
+; CHECK-LABEL: typedevent:
+; CHECK-NEXT:     .Lfunc_begin2:
+; CHECK-NEXT:       .cfi_startproc
+; CHECK-NEXT:       // %bb.0:
+; CHECK-NEXT:       .p2align 2
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK:            cbnz
+; CHECK-SAME:           [[FALLTHROUGH_3:.LBB[0-9_]+]]
+; CHECK-NEXT:       b
+; CHECK-SAME:           [[OUT_OF_RANGE_3:.LBB[0-9_]+]]
+; CHECK-NEXT:     [[FALLTHROUGH_3]]:
+; CHECK-SAME:                        // %end1
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK-NEXT:         Begin XRay typed event
+; CHECK:            bl      __xray_TypedEvent
+; CHECK:              End XRay typed event
+; CHECK-NEXT:     [[OUT_OF_RANGE_3]]:
+; CHECK-SAME:                        // %end2
+; CHECK-NEXT:     .Lxray_sled_{{[0-9]+}}:
+; CHECK-NEXT:         Begin XRay typed event
+; CHECK:            bl      __xray_TypedEvent
+; CHECK:              End XRay typed event
+; CHECK:          .Ltmp
+; CHECK-NEXT:       ret
+entry:
+  br i1 %0, label %end1, label %end2
+
+end1:
+  call void @llvm.xray.typedevent(i64 %type, ptr %event, i64 %size)
+  br label %end2
+  
+end2:
+  tail call void @llvm.xray.typedevent(i64 %size, ptr %event, i64 %type)
+  ret void
+}
+
+declare void @llvm.xray.customevent(ptr, i64)
+declare void @llvm.xray.typedevent(i64, ptr, i64)
 declare i32 @bar()
 declare i32 @baz()
 No newline at end of file