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/CodeGen/MachineBasicBlock.h"
19#include "llvm/CodeGen/MachineDominators.h"
20#include "llvm/CodeGen/MachineFunction.h"
21#include "llvm/CodeGen/MachineFunctionPass.h"
22#include "llvm/CodeGen/MachineInstrBuilder.h"
23#include "llvm/CodeGen/MachineLoopInfo.h"
24#include "llvm/CodeGen/TargetInstrInfo.h"
25#include "llvm/CodeGen/TargetSubtargetInfo.h"
26#include "llvm/IR/Attributes.h"
27#include "llvm/IR/Function.h"
28#include "llvm/InitializePasses.h"
29#include "llvm/Pass.h"
30#include "llvm/Target/TargetMachine.h"
31#include "llvm/TargetParser/Triple.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(Inst: 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(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc))
110 .addImm(Val: T.getOpcode());
111 for (auto &MO : T.operands())
112 MIB.add(MO);
113 Terminators.push_back(Elt: &T);
114 if (T.shouldUpdateCallSiteInfo())
115 MF.eraseCallSiteInfo(MI: &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(Inst: 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(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc));
141 }
142 }
143}
144
145bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
146 auto &F = MF.getFunction();
147 auto InstrAttr = F.getFnAttribute(Kind: "function-instrument");
148 bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
149 InstrAttr.getValueAsString() == "xray-always";
150 bool NeverInstrument = InstrAttr.isStringAttribute() &&
151 InstrAttr.getValueAsString() == "xray-never";
152 if (NeverInstrument && !AlwaysInstrument)
153 return false;
154 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
155
156 uint64_t XRayThreshold = 0;
157 if (!AlwaysInstrument) {
158 bool IgnoreLoops = IgnoreLoopsAttr.isValid();
159 XRayThreshold = F.getFnAttributeAsParsedInteger(
160 Kind: "xray-instruction-threshold", Default: std::numeric_limits<uint64_t>::max());
161 if (XRayThreshold == std::numeric_limits<uint64_t>::max())
162 return false;
163
164 // Count the number of MachineInstr`s in MachineFunction
165 uint64_t MICount = 0;
166 for (const auto &MBB : MF)
167 MICount += MBB.size();
168
169 bool TooFewInstrs = MICount < XRayThreshold;
170
171 if (!IgnoreLoops) {
172 // Get MachineDominatorTree or compute it on the fly if it's unavailable
173 auto *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
174 MachineDominatorTree ComputedMDT;
175 if (!MDT) {
176 ComputedMDT.getBase().recalculate(Func&: MF);
177 MDT = &ComputedMDT;
178 }
179
180 // Get MachineLoopInfo or compute it on the fly if it's unavailable
181 auto *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
182 MachineLoopInfo ComputedMLI;
183 if (!MLI) {
184 ComputedMLI.getBase().analyze(DomTree: MDT->getBase());
185 MLI = &ComputedMLI;
186 }
187
188 // Check if we have a loop.
189 // FIXME: Maybe make this smarter, and see whether the loops are dependent
190 // on inputs or side-effects?
191 if (MLI->empty() && TooFewInstrs)
192 return false; // Function is too small and has no loops.
193 } else if (TooFewInstrs) {
194 // Function is too small
195 return false;
196 }
197 }
198
199 // We look for the first non-empty MachineBasicBlock, so that we can insert
200 // the function instrumentation in the appropriate place.
201 auto MBI = llvm::find_if(
202 Range&: MF, P: [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
203 if (MBI == MF.end())
204 return false; // The function is empty.
205
206 auto *TII = MF.getSubtarget().getInstrInfo();
207 auto &FirstMBB = *MBI;
208 auto &FirstMI = *FirstMBB.begin();
209
210 if (!MF.getSubtarget().isXRaySupported()) {
211 FirstMI.emitError(Msg: "An attempt to perform XRay instrumentation for an"
212 " unsupported target.");
213 return false;
214 }
215
216 if (!F.hasFnAttribute(Kind: "xray-skip-entry")) {
217 // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
218 // MachineFunction.
219 BuildMI(BB&: FirstMBB, I&: FirstMI, MIMD: FirstMI.getDebugLoc(),
220 MCID: TII->get(Opcode: TargetOpcode::PATCHABLE_FUNCTION_ENTER));
221 }
222
223 if (!F.hasFnAttribute(Kind: "xray-skip-exit")) {
224 switch (MF.getTarget().getTargetTriple().getArch()) {
225 case Triple::ArchType::arm:
226 case Triple::ArchType::thumb:
227 case Triple::ArchType::aarch64:
228 case Triple::ArchType::hexagon:
229 case Triple::ArchType::loongarch64:
230 case Triple::ArchType::mips:
231 case Triple::ArchType::mipsel:
232 case Triple::ArchType::mips64:
233 case Triple::ArchType::mips64el: {
234 // For the architectures which don't have a single return instruction
235 InstrumentationOptions op;
236 op.HandleTailcall = false;
237 op.HandleAllReturns = true;
238 prependRetWithPatchableExit(MF, TII, op);
239 break;
240 }
241 case Triple::ArchType::ppc64le: {
242 // PPC has conditional returns. Turn them into branch and plain returns.
243 InstrumentationOptions op;
244 op.HandleTailcall = false;
245 op.HandleAllReturns = true;
246 replaceRetWithPatchableRet(MF, TII, op);
247 break;
248 }
249 default: {
250 // For the architectures that have a single return instruction (such as
251 // RETQ on x86_64).
252 InstrumentationOptions op;
253 op.HandleTailcall = true;
254 op.HandleAllReturns = false;
255 replaceRetWithPatchableRet(MF, TII, op);
256 break;
257 }
258 }
259 }
260 return true;
261}
262
263char XRayInstrumentation::ID = 0;
264char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
265INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
266 "Insert XRay ops", false, false)
267INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
268INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
269 "Insert XRay ops", false, false)
270

source code of llvm/lib/CodeGen/XRayInstrumentation.cpp