[XRay][AArch64] Suppport __xray_customevent/__xray_typedevent

`__xray_customevent` and `__xray_typedevent` are built-in functions in Clang.
With -fxray-instrument, they are lowered to intrinsics llvm.xray.customevent and
llvm.xray.typedevent, respectively. These intrinsics are then lowered to
TargetOpcode::{PATCHABLE_EVENT_CALL,PATCHABLE_TYPED_EVENT_CALL}. The target is
responsible for generating a code sequence that calls either
`__xray_CustomEvent` (with 2 arguments) or `__xray_TypedEvent` (with 3
arguments).

Before patching, the code sequence is prefixed by a branch instruction that
skips the rest of the code sequence. After patching
(compiler-rt/lib/xray/xray_AArch64.cpp), the branch instruction becomes a NOP
and the function call will take effects.

This patch implements the lowering process for
{PATCHABLE_EVENT_CALL,PATCHABLE_TYPED_EVENT_CALL} and implements the runtime.

```
// Lowering of PATCHABLE_EVENT_CALL
.Lxray_sled_N:
  b  #24
  stp x0, x1, [sp, #-16]!
  x0 = reg of op0
  x1 = reg of op1
  bl __xray_CustomEvent
  ldrp x0, x1, [sp], #16
```

As a result, two updated tests in compiler-rt/test/xray/TestCases/Posix/ now
pass on AArch64.

Reviewed By: peter.smith

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

Author: MaskRay

compiler-rt/lib/xray/xray_AArch64.cpp

@@ -105,15 +105,37 @@ bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
   return patchSled(Enable, FuncId, Sled, __xray_FunctionTailExit);
 }
 
+// AArch64AsmPrinter::LowerPATCHABLE_EVENT_CALL generates this code sequence:
+//
+// .Lxray_event_sled_N:
+//   b 1f
+//   save x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+//   set up x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+//   bl __xray_CustomEvent or __xray_TypedEvent
+//   restore x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+// 1f
+//
+// There are 6 instructions for EVENT_CALL and 9 for TYPED_EVENT_CALL.
+//
+// Enable: b .+24 => nop
+// Disable: nop => b .+24
 bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
-                      const XRaySledEntry &Sled)
-    XRAY_NEVER_INSTRUMENT { // FIXME: Implement in aarch64?
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  uint32_t Inst = Enable ? 0xd503201f : 0x14000006;
+  std::atomic_store_explicit(
+      reinterpret_cast<std::atomic<uint32_t> *>(Sled.address()), Inst,
+      std::memory_order_release);
   return false;
 }
 
+// Enable: b +36 => nop
+// Disable: nop => b +36
 bool patchTypedEvent(const bool Enable, const uint32_t FuncId,
                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
-  // FIXME: Implement in aarch64?
+  uint32_t Inst = Enable ? 0xd503201f : 0x14000009;
+  std::atomic_store_explicit(
+      reinterpret_cast<std::atomic<uint32_t> *>(Sled.address()), Inst,
+      std::memory_order_release);
   return false;
 }
 

compiler-rt/lib/xray/xray_trampoline_AArch64.S

@@ -103,4 +103,32 @@ ASM_SYMBOL(__xray_FunctionTailExit):
   RESTORE_REGISTERS
   ret
 
+.global ASM_SYMBOL(__xray_CustomEvent)
+ASM_HIDDEN(__xray_CustomEvent)
+ASM_TYPE_FUNCTION(__xray_CustomEvent)
+ASM_SYMBOL(__xray_CustomEvent):
+  SAVE_REGISTERS
+  adrp x8, ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)
+  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)]
+  cbz x8, 1f
+  blr x8
+1:
+  RESTORE_REGISTERS
+  ret
+ASM_SIZE(__xray_CustomEvent)
+
+.global ASM_SYMBOL(__xray_TypedEvent)
+ASM_HIDDEN(__xray_TypedEvent)
+ASM_TYPE_FUNCTION(__xray_TypedEvent)
+ASM_SYMBOL(__xray_TypedEvent):
+  SAVE_REGISTERS
+  adrp x8, ASM_SYMBOL(_ZN6__xray21XRayPatchedTypedEventE)
+  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray21XRayPatchedTypedEventE)]
+  cbz x8, 1f
+  blr x8
+1:
+  RESTORE_REGISTERS
+  ret
+ASM_SIZE(__xray_TypedEvent)
+
 NO_EXEC_STACK_DIRECTIVE

compiler-rt/test/xray/TestCases/Posix/custom-event-logging.cpp

@@ -5,7 +5,7 @@
 // RUN: %clangxx_xray -std=c++11 -fpic -fpie %s -o %t
 // RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_logfile_base=custom-event-logging.xray-" %run %t 2>&1 | FileCheck %s
 // FIXME: Support this in non-x86_64 as well
-// REQUIRES: x86_64-linux
+// REQUIRES: target={{(aarch64|x86_64)-.*linux.*}}
 // REQUIRES: built-in-llvm-tree
 #include <cstdio>
 #include "xray/xray_interface.h"

compiler-rt/test/xray/TestCases/Posix/typed-event-logging.cpp

@@ -1,7 +1,7 @@
 // RUN: %clangxx_xray %s -o %t
 // RUN: XRAY_OPTIONS=patch_premain=false:verbosity=1 %run %t 2>&1 | FileCheck %s
 
-// REQUIRES: target={{x86_64-.*linux.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*linux.*}}
 
 #include <assert.h>
 #include <stdio.h>

llvm/lib/CodeGen/SelectionDAG/FastISel.cpp

@@ -893,7 +893,7 @@ bool FastISel::selectPatchpoint(const CallInst *I) {
 
 bool FastISel::selectXRayCustomEvent(const CallInst *I) {
   const auto &Triple = TM.getTargetTriple();
-  if (Triple.getArch() != Triple::x86_64 || !Triple.isOSLinux())
+  if (Triple.isAArch64(64) && Triple.getArch() != Triple::x86_64)
     return true; // don't do anything to this instruction.
   SmallVector<MachineOperand, 8> Ops;
   Ops.push_back(MachineOperand::CreateReg(getRegForValue(I->getArgOperand(0)),
@@ -912,7 +912,7 @@ bool FastISel::selectXRayCustomEvent(const CallInst *I) {
 
 bool FastISel::selectXRayTypedEvent(const CallInst *I) {
   const auto &Triple = TM.getTargetTriple();
-  if (Triple.getArch() != Triple::x86_64 || !Triple.isOSLinux())
+  if (Triple.isAArch64(64) && Triple.getArch() != Triple::x86_64)
     return true; // don't do anything to this instruction.
   SmallVector<MachineOperand, 8> Ops;
   Ops.push_back(MachineOperand::CreateReg(getRegForValue(I->getArgOperand(0)),

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -7204,10 +7204,9 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
   }
   case Intrinsic::xray_customevent: {
     // Here we want to make sure that the intrinsic behaves as if it has a
-    // specific calling convention, and only for x86_64.
-    // FIXME: Support other platforms later.
+    // specific calling convention.
     const auto &Triple = DAG.getTarget().getTargetTriple();
-    if (Triple.getArch() != Triple::x86_64)
+    if (!Triple.isAArch64(64) && Triple.getArch() != Triple::x86_64)
       return;
 
     SmallVector<SDValue, 8> Ops;
@@ -7234,10 +7233,9 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
   }
   case Intrinsic::xray_typedevent: {
     // Here we want to make sure that the intrinsic behaves as if it has a
-    // specific calling convention, and only for x86_64.
-    // FIXME: Support other platforms later.
+    // specific calling convention.
     const auto &Triple = DAG.getTarget().getTargetTriple();
-    if (Triple.getArch() != Triple::x86_64)
+    if (!Triple.isAArch64(64) && Triple.getArch() != Triple::x86_64)
       return;
 
     SmallVector<SDValue, 8> Ops;

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -107,6 +107,7 @@ class AArch64AsmPrinter : public AsmPrinter {
   void LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI);
   void LowerPATCHABLE_FUNCTION_EXIT(const MachineInstr &MI);
   void LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI);
+  void LowerPATCHABLE_EVENT_CALL(const MachineInstr &MI, bool Typed);
 
   typedef std::tuple<unsigned, bool, uint32_t> HwasanMemaccessTuple;
   std::map<HwasanMemaccessTuple, MCSymbol *> HwasanMemaccessSymbols;
@@ -323,6 +324,100 @@ void AArch64AsmPrinter::emitSled(const MachineInstr &MI, SledKind Kind) {
   recordSled(CurSled, MI, Kind, 2);
 }
 
+// Emit the following code for Intrinsic::{xray_customevent,xray_typedevent}
+// (built-in functions __xray_customevent/__xray_typedevent).
+//
+// .Lxray_event_sled_N:
+//   b 1f
+//   save x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+//   set up x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+//   bl __xray_CustomEvent or __xray_TypedEvent
+//   restore x0 and x1 (and also x2 for TYPED_EVENT_CALL)
+// 1:
+//
+// There are 6 instructions for EVENT_CALL and 9 for TYPED_EVENT_CALL.
+//
+// Then record a sled of kind CUSTOM_EVENT or TYPED_EVENT.
+// After patching, b .+N will become a nop.
+void AArch64AsmPrinter::LowerPATCHABLE_EVENT_CALL(const MachineInstr &MI,
+                                                  bool Typed) {
+  auto &O = *OutStreamer;
+  MCSymbol *CurSled = OutContext.createTempSymbol("xray_sled_", true);
+  O.emitLabel(CurSled);
+  MCInst MovX0Op0 = MCInstBuilder(AArch64::ORRXrs)
+                        .addReg(AArch64::X0)
+                        .addReg(AArch64::XZR)
+                        .addReg(MI.getOperand(0).getReg())
+                        .addImm(0);
+  MCInst MovX1Op1 = MCInstBuilder(AArch64::ORRXrs)
+                        .addReg(AArch64::X1)
+                        .addReg(AArch64::XZR)
+                        .addReg(MI.getOperand(1).getReg())
+                        .addImm(0);
+  bool MachO = TM.getTargetTriple().isOSBinFormatMachO();
+  auto *Sym = MCSymbolRefExpr::create(
+      OutContext.getOrCreateSymbol(
+          Twine(MachO ? "_" : "") +
+          (Typed ? "__xray_TypedEvent" : "__xray_CustomEvent")),
+      OutContext);
+  if (Typed) {
+    O.AddComment("Begin XRay typed event");
+    EmitToStreamer(O, MCInstBuilder(AArch64::B).addImm(9));
+    EmitToStreamer(O, MCInstBuilder(AArch64::STPXpre)
+                          .addReg(AArch64::SP)
+                          .addReg(AArch64::X0)
+                          .addReg(AArch64::X1)
+                          .addReg(AArch64::SP)
+                          .addImm(-4));
+    EmitToStreamer(O, MCInstBuilder(AArch64::STRXui)
+                          .addReg(AArch64::X2)
+                          .addReg(AArch64::SP)
+                          .addImm(2));
+    EmitToStreamer(O, MovX0Op0);
+    EmitToStreamer(O, MovX1Op1);
+    EmitToStreamer(O, MCInstBuilder(AArch64::ORRXrs)
+                          .addReg(AArch64::X2)
+                          .addReg(AArch64::XZR)
+                          .addReg(MI.getOperand(2).getReg())
+                          .addImm(0));
+    EmitToStreamer(O, MCInstBuilder(AArch64::BL).addExpr(Sym));
+    EmitToStreamer(O, MCInstBuilder(AArch64::LDRXui)
+                          .addReg(AArch64::X2)
+                          .addReg(AArch64::SP)
+                          .addImm(2));
+    O.AddComment("End XRay typed event");
+    EmitToStreamer(O, MCInstBuilder(AArch64::LDPXpost)
+                          .addReg(AArch64::SP)
+                          .addReg(AArch64::X0)
+                          .addReg(AArch64::X1)
+                          .addReg(AArch64::SP)
+                          .addImm(4));
+
+    recordSled(CurSled, MI, SledKind::TYPED_EVENT, 2);
+  } else {
+    O.AddComment("Begin XRay custom event");
+    EmitToStreamer(O, MCInstBuilder(AArch64::B).addImm(6));
+    EmitToStreamer(O, MCInstBuilder(AArch64::STPXpre)
+                          .addReg(AArch64::SP)
+                          .addReg(AArch64::X0)
+                          .addReg(AArch64::X1)
+                          .addReg(AArch64::SP)
+                          .addImm(-2));
+    EmitToStreamer(O, MovX0Op0);
+    EmitToStreamer(O, MovX1Op1);
+    EmitToStreamer(O, MCInstBuilder(AArch64::BL).addExpr(Sym));
+    O.AddComment("End XRay custom event");
+    EmitToStreamer(O, MCInstBuilder(AArch64::LDPXpost)
+                          .addReg(AArch64::SP)
+                          .addReg(AArch64::X0)
+                          .addReg(AArch64::X1)
+                          .addReg(AArch64::SP)
+                          .addImm(2));
+
+    recordSled(CurSled, MI, SledKind::CUSTOM_EVENT, 2);
+  }
+}
+
 void AArch64AsmPrinter::LowerKCFI_CHECK(const MachineInstr &MI) {
   Register AddrReg = MI.getOperand(0).getReg();
   assert(std::next(MI.getIterator())->isCall() &&
@@ -1552,6 +1647,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
   case TargetOpcode::PATCHABLE_TAIL_CALL:
     LowerPATCHABLE_TAIL_CALL(*MI);
     return;
+  case TargetOpcode::PATCHABLE_EVENT_CALL:
+    return LowerPATCHABLE_EVENT_CALL(*MI, false);
+  case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
+    return LowerPATCHABLE_EVENT_CALL(*MI, true);
 
   case AArch64::KCFI_CHECK:
     LowerKCFI_CHECK(*MI);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2757,6 +2757,10 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
   case TargetOpcode::PATCHPOINT:
     return emitPatchPoint(MI, BB);
 
+  case TargetOpcode::PATCHABLE_EVENT_CALL:
+  case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
+    return BB;
+
   case AArch64::CATCHRET:
     return EmitLoweredCatchRet(MI, BB);
   case AArch64::LD1_MXIPXX_H_PSEUDO_B: