[XRay] Support AArch64 in LLVM

[rSerge]

Summary

This patch adds XRay support in LLVM for AArch64 targets.
This patch is one of a series:

• Clang: https://reviews.llvm.org/D26415
• compiler-rt: https://reviews.llvm.org/D26413

[dberris]

lib/Target/AArch64/AArch64AsmPrinter.cpp

159 ↗

(On Diff #77234)

I distinctly remember having to do something special for non-COMDAT sections, especially when doing section-per-function in x86. I suspect you'd want to replicate that logic here somehow, with a //TODO to merge the logic for ELF XRay sleds at a higher level for better re-use.

210–212 ↗

(On Diff #77234)

Probably remove '{}' for single line bodies?

lib/Target/AArch64/AArch64Subtarget.h

176 ↗

(On Diff #77234)

I don't think you need the virtual here especially since you mark it override anyway.

[rSerge]

lib/Target/AArch64/AArch64AsmPrinter.cpp

159 ↗

(On Diff #77234)

Yes, I see it now. Shall this code be CPU-independent? If so, let me look if I can refactor it in this patch...
The CPU-independent code seems to be

  if (Sleds.empty())
    return;
  if (Subtarget->isTargetELF()) {
    auto PrevSection = OutStreamer->getCurrentSectionOnly();
    auto Fn = MF->getFunction();
    MCSection *Section = nullptr;
    if (Fn->hasComdat()) {
      Section = OutContext.getELFSection("xray_instr_map", ELF::SHT_PROGBITS,
                                         ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
                                         Fn->getComdat()->getName());
    } else {
      Section = OutContext.getELFSection("xray_instr_map", ELF::SHT_PROGBITS,
                                         ELF::SHF_ALLOC);
    }

    // Before we switch over, we force a reference to a label inside the
    // xray_instr_map section. Since EmitXRayTable() is always called just
    // before the function's end, we assume that this is happening after the
    // last return instruction.
    //
    // We then align the reference to 16 byte boundaries, which we determined
    // experimentally to be beneficial to avoid causing decoder stalls.
    MCSymbol *Tmp = OutContext.createTempSymbol("xray_synthetic_", true);
    OutStreamer->EmitCodeAlignment(16);
    OutStreamer->EmitSymbolValue(Tmp, 8, false);
    OutStreamer->SwitchSection(Section);
    OutStreamer->EmitLabel(Tmp);
    for (const auto &Sled : Sleds) {
         // CPU-dependent code block here (depends of 64- or 32- bit word width)
    }
    OutStreamer->SwitchSection(PrevSection);
  }
  Sleds.clear();
}

210–212 ↗

(On Diff #77234)

Changing.

lib/Target/AArch64/AArch64Subtarget.h

176 ↗

(On Diff #77234)

Changing.

[dberris]

lib/Target/AArch64/AArch64AsmPrinter.cpp

159 ↗

(On Diff #77234)

Yep, I think you found it. Would be happy to wait for the refactoring to make this simpler moving forward when we cover more architectures.

[amehsan]

lib/Target/AArch64/AArch64AsmPrinter.cpp

155 ↗

(On Diff #77541)

extra space in the beginning can be removed.

[rSerge]

lib/Target/AArch64/AArch64AsmPrinter.cpp

159 ↗

(On Diff #77234)

My team has suggested to defer the refactorring to a separate patch, so to keep the current patch AArch64-only. That refactoring would affect multiple CPUs, but the changes would have in common that they are all for ELF.
So I've uploaded a new version with code duplication and a TODO for the refactoring.

[dberris]

Just a couple style issues, thanks @rSerge -- I'll look forward to the refactoring afterwards to consolidate some of the ELF-specific XRay instrumentation stuff potentially at a higher level at a later time.

lib/Target/AArch64/AArch64AsmPrinter.cpp

162–170 ↗

(On Diff #77541)

Since these are single-statement bodies for the true and false branches, LLVM style seems to prefer not having curlies around them:

if (...)
  Section = ...;
else
  Section = ...;

Also, if you do need compound statements in the branches, the else should be on the same line as the previous }. So:

if (...) {
  ...
} else {
  ...
}

685–686 ↗

(On Diff #77541)

To keep it consistent with surrounding style, maybe add an empty line between the return; and the case ...:? i.e.

case TargetOpcode::...:
  ...
  return;

case TargetOpcode::...;
  ...
  return;

[rSerge]

lib/Target/AArch64/AArch64AsmPrinter.cpp

155 ↗

(On Diff #77541)

Leaving 1 space after //

162–170 ↗

(On Diff #77541)

Thanks for explaining. It's copy-pasted from your code =)
Changing here...

685–686 ↗

(On Diff #77541)

Changing.

[rSerge]

Rebased to the latest version of LLVM sources.

[dberris]

Now that this has been landed (just waiting for the build bots to finish going through the tests) it would be great if we also had tests for aarch64 and maybe arm7 specifically to look for the assembler patterns we want similar to the x86_64 tests.

[rengolin]

In D26412#598292, @dberris wrote:

Now that this has been landed (just waiting for the build bots to finish going through the tests) it would be great if we also had tests for aarch64 and maybe arm7 specifically to look for the assembler patterns we want similar to the x86_64 tests.

Yes, please! :)

[rSerge]

Oops, I remember that I implemented a test for this, but got lost it somewhere... looking...

[rSerge]

No, it was a test for tail calls.
Here is a test and a fix for the jump instruction emitted in sleds: https://reviews.llvm.org/D26805

llvm/trunk/lib/CodeGen/XRayInstrumentation.cpp

bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   // MachineFunction.
   auto *TII = MF.getSubtarget().getInstrInfo();
   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 
   switch (MF.getTarget().getTargetTriple().getArch()) {
   case Triple::ArchType::arm:
   case Triple::ArchType::thumb:
+  case Triple::ArchType::aarch64:
     // For the architectures which don't have a single return instruction
     prependRetWithPatchableExit(MF, TII);
     break;
   default:
     // For the architectures that have a single return instruction (such as
     //   RETQ on x86_64).
     replaceRetWithPatchableRet(MF, TII);
     break;
   }
   return true;
 }
 
 char XRayInstrumentation::ID = 0;
 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
 INITIALIZE_PASS(XRayInstrumentation, "xray-instrumentation", "Insert XRay ops",
                 false, false)

llvm/trunk/lib/Target/AArch64/AArch64AsmPrinter.cpp

 #include "llvm/IR/DebugInfo.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstBuilder.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCSymbolELF.h"
+#include "llvm/MC/MCSectionELF.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "asm-printer"
 
 namespace {
public:
   bool lowerOperand(const MachineOperand &MO, MCOperand &MCOp) const {
     return MCInstLowering.lowerOperand(MO, MCOp);
   }
 
   void LowerSTACKMAP(MCStreamer &OutStreamer, StackMaps &SM,
                      const MachineInstr &MI);
   void LowerPATCHPOINT(MCStreamer &OutStreamer, StackMaps &SM,
                        const MachineInstr &MI);
+
+  void LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI);
+  void LowerPATCHABLE_FUNCTION_EXIT(const MachineInstr &MI);
+  void LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI);
+
+  void EmitXRayTable();
+  void EmitSled(const MachineInstr &MI, SledKind Kind);
+
   /// \brief tblgen'erated driver function for lowering simple MI->MC
   /// pseudo instructions.
   bool emitPseudoExpansionLowering(MCStreamer &OutStreamer,
                                    const MachineInstr *MI);
 
   void EmitInstruction(const MachineInstr *MI) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AsmPrinter::getAnalysisUsage(AU);
     AU.setPreservesAll();
   }
 
   bool runOnMachineFunction(MachineFunction &F) override {
     AArch64FI = F.getInfo<AArch64FunctionInfo>();
     STI = static_cast<const AArch64Subtarget*>(&F.getSubtarget());
-    return AsmPrinter::runOnMachineFunction(F);
+    bool Result = AsmPrinter::runOnMachineFunction(F);
+    EmitXRayTable();
+    return Result;
   }
 
 private:
   void printOperand(const MachineInstr *MI, unsigned OpNum, raw_ostream &O);
   bool printAsmMRegister(const MachineOperand &MO, char Mode, raw_ostream &O);
   bool printAsmRegInClass(const MachineOperand &MO,
                           const TargetRegisterClass *RC, bool isVector,
                           raw_ostream &O);
private:
   typedef std::map<const MachineInstr *, MCSymbol *> MInstToMCSymbol;
   MInstToMCSymbol LOHInstToLabel;
 };
 
 } // end of anonymous namespace
 
 //===----------------------------------------------------------------------===//
 
+void AArch64AsmPrinter::LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI)
+{
+  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)
+{
+  EmitSled(MI, SledKind::TAIL_CALL);
+}
+
+void AArch64AsmPrinter::EmitXRayTable()
+{
+  //TODO: merge the logic for ELF XRay sleds at a higher level, so to avoid
+  // code duplication as it is now for x86_64, ARM32 and AArch64.
+  if (Sleds.empty())
+    return;
+  if (STI->isTargetELF()) {
+    auto PrevSection = OutStreamer->getCurrentSectionOnly();
+    auto Fn = MF->getFunction();
+    MCSection *Section;
+    if (Fn->hasComdat())
+      Section = OutContext.getELFSection("xray_instr_map", ELF::SHT_PROGBITS,
+        ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
+        Fn->getComdat()->getName());
+    else
+      Section = OutContext.getELFSection("xray_instr_map", ELF::SHT_PROGBITS,
+        ELF::SHF_ALLOC);
+
+    // Before we switch over, we force a reference to a label inside the
+    // xray_instr_map section. Since EmitXRayTable() is always called just
+    // before the function's end, we assume that this is happening after the
+    // last return instruction.
+    //
+    // We then align the reference to 16 byte boundaries, which we determined
+    // experimentally to be beneficial to avoid causing decoder stalls.
+    MCSymbol *Tmp = OutContext.createTempSymbol("xray_synthetic_", true);
+    OutStreamer->EmitCodeAlignment(16);
+    OutStreamer->EmitSymbolValue(Tmp, 8, false);
+    OutStreamer->SwitchSection(Section);
+    OutStreamer->EmitLabel(Tmp);
+    for (const auto &Sled : Sleds) {
+      OutStreamer->EmitSymbolValue(Sled.Sled, 8);
+      OutStreamer->EmitSymbolValue(CurrentFnSym, 8);
+      auto Kind = static_cast<uint8_t>(Sled.Kind);
+      OutStreamer->EmitBytes(
+        StringRef(reinterpret_cast<const char *>(&Kind), 1));
+      OutStreamer->EmitBytes(
+        StringRef(reinterpret_cast<const char *>(&Sled.AlwaysInstrument), 1));
+      OutStreamer->EmitZeros(14);
+    }
+    OutStreamer->SwitchSection(PrevSection);
+  }
+  Sleds.clear();
+}
+
+void AArch64AsmPrinter::EmitSled(const MachineInstr &MI, SledKind Kind)
+{
+  static const int8_t NoopsInSledCount = 7;
+  // We want to emit the following pattern:
+  //
+  // .Lxray_sled_N:
+  //   ALIGN
+  //   B #32
+  //   ; 7 NOP instructions (28 bytes)
+  // .tmpN
+  //
+  // We need the 28 bytes (7 instructions) because at runtime, we'd be patching
+  // over the full 32 bytes (8 instructions) with the following pattern:
+  //
+  //   STP X0, X30, [SP, #-16]! ; push X0 and the link register to the stack
+  //   LDR W0, #12 ; W0 := function ID
+  //   LDR X16,#12 ; X16 := addr of __xray_FunctionEntry or __xray_FunctionExit
+  //   BLR X16 ; call the tracing trampoline
+  //   ;DATA: 32 bits of function ID
+  //   ;DATA: lower 32 bits of the address of the trampoline
+  //   ;DATA: higher 32 bits of the address of the trampoline
+  //   LDP X0, X30, [SP], #16 ; pop X0 and the link register from the stack
+  //
+  OutStreamer->EmitCodeAlignment(4);
+  auto CurSled = OutContext.createTempSymbol("xray_sled_", true);
+  OutStreamer->EmitLabel(CurSled);
+  auto Target = OutContext.createTempSymbol();
+
+  // Emit "B #32" instruction, which jumps over the next 28 bytes.
+  EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::B).addImm(32));
+
+  for (int8_t I = 0; I < NoopsInSledCount; I++)
+    EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::HINT).addImm(0));
+
+  OutStreamer->EmitLabel(Target);
+  recordSled(CurSled, MI, Kind);
+}
+
 void AArch64AsmPrinter::EmitEndOfAsmFile(Module &M) {
   const Triple &TT = TM.getTargetTriple();
   if (TT.isOSBinFormatMachO()) {
     // Funny Darwin hack: This flag tells the linker that no global symbols
     // contain code that falls through to other global symbols (e.g. the obvious
     // implementation of multiple entry points).  If this doesn't occur, the
     // linker can safely perform dead code stripping.  Since LLVM never
     // generates code that does this, it is always safe to set.
case AArch64::FMOVD0:
     EmitFMov0(*MI);
     return;
 
   case TargetOpcode::STACKMAP:
     return LowerSTACKMAP(*OutStreamer, SM, *MI);
 
   case TargetOpcode::PATCHPOINT:
     return LowerPATCHPOINT(*OutStreamer, SM, *MI);
+
+  case TargetOpcode::PATCHABLE_FUNCTION_ENTER:
+    LowerPATCHABLE_FUNCTION_ENTER(*MI);
+    return;
+
+  case TargetOpcode::PATCHABLE_FUNCTION_EXIT:
+    LowerPATCHABLE_FUNCTION_EXIT(*MI);
+    return;
+
+  case TargetOpcode::PATCHABLE_TAIL_CALL:
+    LowerPATCHABLE_TAIL_CALL(*MI);
+    return;
   }
 
   // Finally, do the automated lowerings for everything else.
   MCInst TmpInst;
   MCInstLowering.Lower(MI, TmpInst);
   EmitToStreamer(*OutStreamer, TmpInst);
 }
 
 // Force static initialization.
 extern "C" void LLVMInitializeAArch64AsmPrinter() {
   RegisterAsmPrinter<AArch64AsmPrinter> X(getTheAArch64leTarget());
   RegisterAsmPrinter<AArch64AsmPrinter> Y(getTheAArch64beTarget());
   RegisterAsmPrinter<AArch64AsmPrinter> Z(getTheARM64Target());
 }

llvm/trunk/lib/Target/AArch64/AArch64Subtarget.h

public:
   bool hasV8_2aOps() const { return HasV8_2aOps; }
 
   bool hasZeroCycleRegMove() const { return HasZeroCycleRegMove; }
 
   bool hasZeroCycleZeroing() const { return HasZeroCycleZeroing; }
 
   bool requiresStrictAlign() const { return StrictAlign; }
 
+  bool isXRaySupported() const override { return true; }
+
   bool isX18Reserved() const { return ReserveX18; }
   bool hasFPARMv8() const { return HasFPARMv8; }
   bool hasNEON() const { return HasNEON; }
   bool hasCrypto() const { return HasCrypto; }
   bool hasCRC() const { return HasCRC; }
   bool hasRAS() const { return HasRAS; }
   bool balanceFPOps() const { return BalanceFPOps; }
   bool predictableSelectIsExpensive() const {