XRayInstrumentation.cpp revision 360784 
1//===- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. --===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements a MachineFunctionPass that inserts the appropriate
10// XRay instrumentation instructions. We look for XRay-specific attributes
11// on the function to determine whether we should insert the replacement
12// operations.
13//
14//===---------------------------------------------------------------------===//
15
16#include "llvm/ADT/STLExtras.h"
17#include "llvm/ADT/SmallVector.h"
18#include "llvm/ADT/Triple.h"
19#include "llvm/CodeGen/MachineBasicBlock.h"
20#include "llvm/CodeGen/MachineDominators.h"
21#include "llvm/CodeGen/MachineFunction.h"
22#include "llvm/CodeGen/MachineFunctionPass.h"
23#include "llvm/CodeGen/MachineInstrBuilder.h"
24#include "llvm/CodeGen/MachineLoopInfo.h"
25#include "llvm/CodeGen/TargetInstrInfo.h"
26#include "llvm/CodeGen/TargetSubtargetInfo.h"
27#include "llvm/IR/Attributes.h"
28#include "llvm/IR/Function.h"
29#include "llvm/InitializePasses.h"
30#include "llvm/Pass.h"
31#include "llvm/Target/TargetMachine.h"
32
33using namespace llvm;
34
35namespace {
36
37struct InstrumentationOptions {
38  // Whether to emit PATCHABLE_TAIL_CALL.
39  bool HandleTailcall;
40
41  // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
42  // return, e.g. conditional return.
43  bool HandleAllReturns;
44};
45
46struct XRayInstrumentation : public MachineFunctionPass {
47  static char ID;
48
49  XRayInstrumentation() : MachineFunctionPass(ID) {
50    initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
51  }
52
53  void getAnalysisUsage(AnalysisUsage &AU) const override {
54    AU.setPreservesCFG();
55    AU.addPreserved<MachineLoopInfo>();
56    AU.addPreserved<MachineDominatorTree>();
57    MachineFunctionPass::getAnalysisUsage(AU);
58  }
59
60  bool runOnMachineFunction(MachineFunction &MF) override;
61
62private:
63  // Replace the original RET instruction with the exit sled code ("patchable
64  //   ret" pseudo-instruction), so that at runtime XRay can replace the sled
65  //   with a code jumping to XRay trampoline, which calls the tracing handler
66  //   and, in the end, issues the RET instruction.
67  // This is the approach to go on CPUs which have a single RET instruction,
68  //   like x86/x86_64.
69  void replaceRetWithPatchableRet(MachineFunction &MF,
70                                  const TargetInstrInfo *TII,
71                                  InstrumentationOptions);
72
73  // Prepend the original return instruction with the exit sled code ("patchable
74  //   function exit" pseudo-instruction), preserving the original return
75  //   instruction just after the exit sled code.
76  // This is the approach to go on CPUs which have multiple options for the
77  //   return instruction, like ARM. For such CPUs we can't just jump into the
78  //   XRay trampoline and issue a single return instruction there. We rather
79  //   have to call the trampoline and return from it to the original return
80  //   instruction of the function being instrumented.
81  void prependRetWithPatchableExit(MachineFunction &MF,
82                                   const TargetInstrInfo *TII,
83                                   InstrumentationOptions);
84};
85
86} // end anonymous namespace
87
88void XRayInstrumentation::replaceRetWithPatchableRet(
89    MachineFunction &MF, const TargetInstrInfo *TII,
90    InstrumentationOptions op) {
91  // We look for *all* terminators and returns, then replace those with
92  // PATCHABLE_RET instructions.
93  SmallVector<MachineInstr *, 4> Terminators;
94  for (auto &MBB : MF) {
95    for (auto &T : MBB.terminators()) {
96      unsigned Opc = 0;
97      if (T.isReturn() &&
98          (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
99        // Replace return instructions with:
100        //   PATCHABLE_RET <Opcode>, <Operand>...
101        Opc = TargetOpcode::PATCHABLE_RET;
102      }
103      if (TII->isTailCall(T) && op.HandleTailcall) {
104        // Treat the tail call as a return instruction, which has a
105        // different-looking sled than the normal return case.
106        Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
107      }
108      if (Opc != 0) {
109        auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
110                       .addImm(T.getOpcode());
111        for (auto &MO : T.operands())
112          MIB.add(MO);
113        Terminators.push_back(&T);
114        if (T.isCall())
115          MF.eraseCallSiteInfo(&T);
116      }
117    }
118  }
119
120  for (auto &I : Terminators)
121    I->eraseFromParent();
122}
123
124void XRayInstrumentation::prependRetWithPatchableExit(
125    MachineFunction &MF, const TargetInstrInfo *TII,
126    InstrumentationOptions op) {
127  for (auto &MBB : MF)
128    for (auto &T : MBB.terminators()) {
129      unsigned Opc = 0;
130      if (T.isReturn() &&
131          (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
132        Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
133      }
134      if (TII->isTailCall(T) && op.HandleTailcall) {
135        Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
136      }
137      if (Opc != 0) {
138        // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
139        //   PATCHABLE_TAIL_CALL .
140        BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
141      }
142    }
143}
144
145bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
146  auto &F = MF.getFunction();
147  auto InstrAttr = F.getFnAttribute("function-instrument");
148  bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
149                          InstrAttr.isStringAttribute() &&
150                          InstrAttr.getValueAsString() == "xray-always";
151  Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
152  unsigned XRayThreshold = 0;
153  if (!AlwaysInstrument) {
154    if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
155      return false; // XRay threshold attribute not found.
156    if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
157      return false; // Invalid value for threshold.
158
159    // Count the number of MachineInstr`s in MachineFunction
160    int64_t MICount = 0;
161    for (const auto &MBB : MF)
162      MICount += MBB.size();
163
164    // Get MachineDominatorTree or compute it on the fly if it's unavailable
165    auto *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
166    MachineDominatorTree ComputedMDT;
167    if (!MDT) {
168      ComputedMDT.getBase().recalculate(MF);
169      MDT = &ComputedMDT;
170    }
171
172    // Get MachineLoopInfo or compute it on the fly if it's unavailable
173    auto *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
174    MachineLoopInfo ComputedMLI;
175    if (!MLI) {
176      ComputedMLI.getBase().analyze(MDT->getBase());
177      MLI = &ComputedMLI;
178    }
179
180    // Check if we have a loop.
181    // FIXME: Maybe make this smarter, and see whether the loops are dependent
182    // on inputs or side-effects?
183    if (MLI->empty() && MICount < XRayThreshold)
184      return false; // Function is too small and has no loops.
185  }
186
187  // We look for the first non-empty MachineBasicBlock, so that we can insert
188  // the function instrumentation in the appropriate place.
189  auto MBI = llvm::find_if(
190      MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
191  if (MBI == MF.end())
192    return false; // The function is empty.
193
194  auto *TII = MF.getSubtarget().getInstrInfo();
195  auto &FirstMBB = *MBI;
196  auto &FirstMI = *FirstMBB.begin();
197
198  if (!MF.getSubtarget().isXRaySupported()) {
199    FirstMI.emitError("An attempt to perform XRay instrumentation for an"
200                      " unsupported target.");
201    return false;
202  }
203
204  // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
205  // MachineFunction.
206  BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
207          TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
208
209  switch (MF.getTarget().getTargetTriple().getArch()) {
210  case Triple::ArchType::arm:
211  case Triple::ArchType::thumb:
212  case Triple::ArchType::aarch64:
213  case Triple::ArchType::mips:
214  case Triple::ArchType::mipsel:
215  case Triple::ArchType::mips64:
216  case Triple::ArchType::mips64el: {
217    // For the architectures which don't have a single return instruction
218    InstrumentationOptions op;
219    op.HandleTailcall = false;
220    op.HandleAllReturns = true;
221    prependRetWithPatchableExit(MF, TII, op);
222    break;
223  }
224  case Triple::ArchType::ppc64le: {
225    // PPC has conditional returns. Turn them into branch and plain returns.
226    InstrumentationOptions op;
227    op.HandleTailcall = false;
228    op.HandleAllReturns = true;
229    replaceRetWithPatchableRet(MF, TII, op);
230    break;
231  }
232  default: {
233    // For the architectures that have a single return instruction (such as
234    //   RETQ on x86_64).
235    InstrumentationOptions op;
236    op.HandleTailcall = true;
237    op.HandleAllReturns = false;
238    replaceRetWithPatchableRet(MF, TII, op);
239    break;
240  }
241  }
242  return true;
243}
244
245char XRayInstrumentation::ID = 0;
246char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
247INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
248                      "Insert XRay ops", false, false)
249INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
250INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
251                    "Insert XRay ops", false, false)
252