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.shouldUpdateCallSiteInfo())
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  auto ThresholdAttr = F.getFnAttribute("xray-instruction-threshold");
152  auto IgnoreLoopsAttr = F.getFnAttribute("xray-ignore-loops");
153  unsigned int XRayThreshold = 0;
154  if (!AlwaysInstrument) {
155    if (ThresholdAttr.hasAttribute(Attribute::None) ||
156        !ThresholdAttr.isStringAttribute())
157      return false; // XRay threshold attribute not found.
158    if (ThresholdAttr.getValueAsString().getAsInteger(10, XRayThreshold))
159      return false; // Invalid value for threshold.
160
161    bool IgnoreLoops = !IgnoreLoopsAttr.hasAttribute(Attribute::None);
162
163    // Count the number of MachineInstr`s in MachineFunction
164    int64_t MICount = 0;
165    for (const auto &MBB : MF)
166      MICount += MBB.size();
167
168    bool TooFewInstrs = MICount < XRayThreshold;
169
170    if (!IgnoreLoops) {
171      // Get MachineDominatorTree or compute it on the fly if it's unavailable
172      auto *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
173      MachineDominatorTree ComputedMDT;
174      if (!MDT) {
175        ComputedMDT.getBase().recalculate(MF);
176        MDT = &ComputedMDT;
177      }
178
179      // Get MachineLoopInfo or compute it on the fly if it's unavailable
180      auto *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
181      MachineLoopInfo ComputedMLI;
182      if (!MLI) {
183        ComputedMLI.getBase().analyze(MDT->getBase());
184        MLI = &ComputedMLI;
185      }
186
187      // Check if we have a loop.
188      // FIXME: Maybe make this smarter, and see whether the loops are dependent
189      // on inputs or side-effects?
190      if (MLI->empty() && TooFewInstrs)
191        return false; // Function is too small and has no loops.
192    } else if (TooFewInstrs) {
193      // Function is too small
194      return false;
195    }
196  }
197
198  // We look for the first non-empty MachineBasicBlock, so that we can insert
199  // the function instrumentation in the appropriate place.
200  auto MBI = llvm::find_if(
201      MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
202  if (MBI == MF.end())
203    return false; // The function is empty.
204
205  auto *TII = MF.getSubtarget().getInstrInfo();
206  auto &FirstMBB = *MBI;
207  auto &FirstMI = *FirstMBB.begin();
208
209  if (!MF.getSubtarget().isXRaySupported()) {
210    FirstMI.emitError("An attempt to perform XRay instrumentation for an"
211                      " unsupported target.");
212    return false;
213  }
214
215  if (!F.hasFnAttribute("xray-skip-entry")) {
216    // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
217    // MachineFunction.
218    BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
219            TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
220  }
221
222  if (!F.hasFnAttribute("xray-skip-exit")) {
223    switch (MF.getTarget().getTargetTriple().getArch()) {
224    case Triple::ArchType::arm:
225    case Triple::ArchType::thumb:
226    case Triple::ArchType::aarch64:
227    case Triple::ArchType::mips:
228    case Triple::ArchType::mipsel:
229    case Triple::ArchType::mips64:
230    case Triple::ArchType::mips64el: {
231      // For the architectures which don't have a single return instruction
232      InstrumentationOptions op;
233      op.HandleTailcall = false;
234      op.HandleAllReturns = true;
235      prependRetWithPatchableExit(MF, TII, op);
236      break;
237    }
238    case Triple::ArchType::ppc64le: {
239      // PPC has conditional returns. Turn them into branch and plain returns.
240      InstrumentationOptions op;
241      op.HandleTailcall = false;
242      op.HandleAllReturns = true;
243      replaceRetWithPatchableRet(MF, TII, op);
244      break;
245    }
246    default: {
247      // For the architectures that have a single return instruction (such as
248      //   RETQ on x86_64).
249      InstrumentationOptions op;
250      op.HandleTailcall = true;
251      op.HandleAllReturns = false;
252      replaceRetWithPatchableRet(MF, TII, op);
253      break;
254    }
255    }
256  }
257  return true;
258}
259
260char XRayInstrumentation::ID = 0;
261char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
262INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
263                      "Insert XRay ops", false, false)
264INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
265INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
266                    "Insert XRay ops", false, false)
267