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/CodeGen/XRayInstrumentation.h"
17#include "llvm/ADT/STLExtras.h"
18#include "llvm/ADT/SmallVector.h"
19#include "llvm/CodeGen/MachineBasicBlock.h"
20#include "llvm/CodeGen/MachineDominators.h"
21#include "llvm/CodeGen/MachineFunction.h"
22#include "llvm/CodeGen/MachineFunctionAnalysis.h"
23#include "llvm/CodeGen/MachineFunctionPass.h"
24#include "llvm/CodeGen/MachineInstrBuilder.h"
25#include "llvm/CodeGen/MachineLoopInfo.h"
26#include "llvm/CodeGen/MachinePassManager.h"
27#include "llvm/CodeGen/TargetInstrInfo.h"
28#include "llvm/CodeGen/TargetSubtargetInfo.h"
29#include "llvm/IR/Attributes.h"
30#include "llvm/IR/DiagnosticInfo.h"
31#include "llvm/IR/Function.h"
32#include "llvm/InitializePasses.h"
33#include "llvm/Pass.h"
34#include "llvm/Target/TargetMachine.h"
35#include "llvm/TargetParser/Triple.h"
36
37using namespace llvm;
38
39namespace {
40
41struct InstrumentationOptions {
42 // Whether to emit PATCHABLE_TAIL_CALL.
43 bool HandleTailcall;
44
45 // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
46 // return, e.g. conditional return.
47 bool HandleAllReturns;
48};
49
50struct XRayInstrumentationLegacy : public MachineFunctionPass {
51 static char ID;
52
53 XRayInstrumentationLegacy() : MachineFunctionPass(ID) {}
54
55 void getAnalysisUsage(AnalysisUsage &AU) const override {
56 AU.setPreservesCFG();
57 MachineFunctionPass::getAnalysisUsage(AU);
58 }
59
60 bool runOnMachineFunction(MachineFunction &MF) override;
61};
62
63struct XRayInstrumentation {
64 XRayInstrumentation(MachineDominatorTree *MDT, MachineLoopInfo *MLI)
65 : MDT(MDT), MLI(MLI) {}
66
67 bool run(MachineFunction &MF);
68
69 // Methods for use in the NPM and legacy passes, can be removed once migration
70 // is complete.
71 static bool alwaysInstrument(Function &F) {
72 auto InstrAttr = F.getFnAttribute(Kind: "function-instrument");
73 return InstrAttr.isStringAttribute() &&
74 InstrAttr.getValueAsString() == "xray-always";
75 }
76
77 static bool needMDTAndMLIAnalyses(Function &F) {
78 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
79 auto AlwaysInstrument = XRayInstrumentation::alwaysInstrument(F);
80 return !AlwaysInstrument && !IgnoreLoopsAttr.isValid();
81 }
82
83private:
84 // Replace the original RET instruction with the exit sled code ("patchable
85 // ret" pseudo-instruction), so that at runtime XRay can replace the sled
86 // with a code jumping to XRay trampoline, which calls the tracing handler
87 // and, in the end, issues the RET instruction.
88 // This is the approach to go on CPUs which have a single RET instruction,
89 // like x86/x86_64.
90 void replaceRetWithPatchableRet(MachineFunction &MF,
91 const TargetInstrInfo *TII,
92 InstrumentationOptions);
93
94 // Prepend the original return instruction with the exit sled code ("patchable
95 // function exit" pseudo-instruction), preserving the original return
96 // instruction just after the exit sled code.
97 // This is the approach to go on CPUs which have multiple options for the
98 // return instruction, like ARM. For such CPUs we can't just jump into the
99 // XRay trampoline and issue a single return instruction there. We rather
100 // have to call the trampoline and return from it to the original return
101 // instruction of the function being instrumented.
102 void prependRetWithPatchableExit(MachineFunction &MF,
103 const TargetInstrInfo *TII,
104 InstrumentationOptions);
105
106 MachineDominatorTree *MDT;
107 MachineLoopInfo *MLI;
108};
109
110} // end anonymous namespace
111
112void XRayInstrumentation::replaceRetWithPatchableRet(
113 MachineFunction &MF, const TargetInstrInfo *TII,
114 InstrumentationOptions op) {
115 // We look for *all* terminators and returns, then replace those with
116 // PATCHABLE_RET instructions.
117 SmallVector<MachineInstr *, 4> Terminators;
118 for (auto &MBB : MF) {
119 for (auto &T : MBB.terminators()) {
120 unsigned Opc = 0;
121 if (T.isReturn() &&
122 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
123 // Replace return instructions with:
124 // PATCHABLE_RET <Opcode>, <Operand>...
125 Opc = TargetOpcode::PATCHABLE_RET;
126 }
127 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
128 // Treat the tail call as a return instruction, which has a
129 // different-looking sled than the normal return case.
130 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
131 }
132 if (Opc != 0) {
133 auto MIB = BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc))
134 .addImm(Val: T.getOpcode());
135 for (auto &MO : T.operands())
136 MIB.add(MO);
137 Terminators.push_back(Elt: &T);
138 if (T.shouldUpdateAdditionalCallInfo())
139 MF.eraseAdditionalCallInfo(MI: &T);
140 }
141 }
142 }
143
144 for (auto &I : Terminators)
145 I->eraseFromParent();
146}
147
148void XRayInstrumentation::prependRetWithPatchableExit(
149 MachineFunction &MF, const TargetInstrInfo *TII,
150 InstrumentationOptions op) {
151 for (auto &MBB : MF)
152 for (auto &T : MBB.terminators()) {
153 unsigned Opc = 0;
154 if (T.isReturn() &&
155 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
156 Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
157 }
158 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
159 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
160 }
161 if (Opc != 0) {
162 // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
163 // PATCHABLE_TAIL_CALL .
164 BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc));
165 }
166 }
167}
168
169PreservedAnalyses
170XRayInstrumentationPass::run(MachineFunction &MF,
171 MachineFunctionAnalysisManager &MFAM) {
172 MachineDominatorTree *MDT = nullptr;
173 MachineLoopInfo *MLI = nullptr;
174
175 if (XRayInstrumentation::needMDTAndMLIAnalyses(F&: MF.getFunction())) {
176 MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(IR&: MF);
177 MLI = MFAM.getCachedResult<MachineLoopAnalysis>(IR&: MF);
178 }
179
180 if (!XRayInstrumentation(MDT, MLI).run(MF))
181 return PreservedAnalyses::all();
182
183 auto PA = getMachineFunctionPassPreservedAnalyses();
184 PA.preserveSet<CFGAnalyses>();
185 return PA;
186}
187
188bool XRayInstrumentationLegacy::runOnMachineFunction(MachineFunction &MF) {
189 MachineDominatorTree *MDT = nullptr;
190 MachineLoopInfo *MLI = nullptr;
191 if (XRayInstrumentation::needMDTAndMLIAnalyses(F&: MF.getFunction())) {
192 auto *MDTWrapper =
193 getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
194 MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
195 auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
196 MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
197 }
198 return XRayInstrumentation(MDT, MLI).run(MF);
199}
200
201bool XRayInstrumentation::run(MachineFunction &MF) {
202 auto &F = MF.getFunction();
203 auto InstrAttr = F.getFnAttribute(Kind: "function-instrument");
204 bool AlwaysInstrument = alwaysInstrument(F);
205 bool NeverInstrument = InstrAttr.isStringAttribute() &&
206 InstrAttr.getValueAsString() == "xray-never";
207 if (NeverInstrument && !AlwaysInstrument)
208 return false;
209 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
210
211 uint64_t XRayThreshold = 0;
212 if (!AlwaysInstrument) {
213 bool IgnoreLoops = IgnoreLoopsAttr.isValid();
214 XRayThreshold = F.getFnAttributeAsParsedInteger(
215 Kind: "xray-instruction-threshold", Default: std::numeric_limits<uint64_t>::max());
216 if (XRayThreshold == std::numeric_limits<uint64_t>::max())
217 return false;
218
219 // Count the number of MachineInstr`s in MachineFunction
220 uint64_t MICount = 0;
221 for (const auto &MBB : MF)
222 MICount += MBB.size();
223
224 bool TooFewInstrs = MICount < XRayThreshold;
225
226 if (!IgnoreLoops) {
227 // Get MachineLoopInfo or compute it on the fly if it's unavailable,
228 // which needs a MachineDominatorTree only for an irreducible CFG.
229 MachineDominatorTree ComputedMDT;
230 MachineLoopInfo ComputedMLI;
231 if (!MLI) {
232 ComputedMLI.calculate(MF, GetDomTree: [&]() -> const MachineDominatorTree & {
233 if (!MDT) {
234 ComputedMDT.recalculate(Func&: MF);
235 MDT = &ComputedMDT;
236 }
237 return *MDT;
238 });
239 MLI = &ComputedMLI;
240 }
241
242 // Check if we have a loop.
243 // FIXME: Maybe make this smarter, and see whether the loops are dependent
244 // on inputs or side-effects?
245 if (MLI->empty() && TooFewInstrs)
246 return false; // Function is too small and has no loops.
247 } else if (TooFewInstrs) {
248 // Function is too small
249 return false;
250 }
251 }
252
253 // We look for the first non-empty MachineBasicBlock, so that we can insert
254 // the function instrumentation in the appropriate place.
255 auto MBI = llvm::find_if(
256 Range&: MF, P: [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
257 if (MBI == MF.end())
258 return false; // The function is empty.
259
260 auto *TII = MF.getSubtarget().getInstrInfo();
261 auto &FirstMBB = *MBI;
262 auto &FirstMI = *FirstMBB.begin();
263
264 if (!MF.getSubtarget().isXRaySupported()) {
265
266 const Function &Fn = FirstMBB.getParent()->getFunction();
267 Fn.getContext().diagnose(DI: DiagnosticInfoUnsupported(
268 Fn, "An attempt to perform XRay instrumentation for an"
269 " unsupported target."));
270
271 return false;
272 }
273
274 if (!F.hasFnAttribute(Kind: "xray-skip-entry")) {
275 // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
276 // MachineFunction.
277 BuildMI(BB&: FirstMBB, I&: FirstMI, MIMD: FirstMI.getDebugLoc(),
278 MCID: TII->get(Opcode: TargetOpcode::PATCHABLE_FUNCTION_ENTER));
279 }
280
281 if (!F.hasFnAttribute(Kind: "xray-skip-exit")) {
282 switch (MF.getTarget().getTargetTriple().getArch()) {
283 case Triple::ArchType::arm:
284 case Triple::ArchType::thumb:
285 case Triple::ArchType::aarch64:
286 case Triple::ArchType::hexagon:
287 case Triple::ArchType::loongarch64:
288 case Triple::ArchType::mips:
289 case Triple::ArchType::mipsel:
290 case Triple::ArchType::mips64:
291 case Triple::ArchType::mips64el:
292 case Triple::ArchType::riscv32:
293 case Triple::ArchType::riscv64: {
294 // For the architectures which don't have a single return instruction
295 InstrumentationOptions op;
296 // AArch64 and RISC-V support patching tail calls.
297 op.HandleTailcall = MF.getTarget().getTargetTriple().isAArch64() ||
298 MF.getTarget().getTargetTriple().isRISCV();
299 op.HandleAllReturns = true;
300 prependRetWithPatchableExit(MF, TII, op);
301 break;
302 }
303 case Triple::ArchType::ppc64le:
304 case Triple::ArchType::systemz: {
305 // PPC has conditional returns. Turn them into branch and plain returns.
306 InstrumentationOptions op;
307 op.HandleTailcall = false;
308 op.HandleAllReturns = true;
309 replaceRetWithPatchableRet(MF, TII, op);
310 break;
311 }
312 default: {
313 // For the architectures that have a single return instruction (such as
314 // RETQ on x86_64).
315 InstrumentationOptions op;
316 op.HandleTailcall = true;
317 op.HandleAllReturns = false;
318 replaceRetWithPatchableRet(MF, TII, op);
319 break;
320 }
321 }
322 }
323 return true;
324}
325
326char XRayInstrumentationLegacy::ID = 0;
327char &llvm::XRayInstrumentationID = XRayInstrumentationLegacy::ID;
328INITIALIZE_PASS_BEGIN(XRayInstrumentationLegacy, "xray-instrumentation",
329 "Insert XRay ops", false, false)
330INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
331INITIALIZE_PASS_END(XRayInstrumentationLegacy, "xray-instrumentation",
332 "Insert XRay ops", false, false)
333